You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
88 lines
2.2 KiB
88 lines
2.2 KiB
# Refer https://github.com/intel-isl/MiDaS |
|
"""Utils for monoDepth. |
|
""" |
|
import sys |
|
import re |
|
import numpy as np |
|
import cv2 |
|
|
|
|
|
def write_pfm(path, image, scale=1): |
|
"""Write pfm file. |
|
|
|
Args: |
|
path (str): pathto file |
|
image (array): data |
|
scale (int, optional): Scale. Defaults to 1. |
|
""" |
|
|
|
with open(path, "wb") as file: |
|
color = None |
|
|
|
if image.dtype.name != "float32": |
|
raise Exception("Image dtype must be float32.") |
|
|
|
image = np.flipud(image) |
|
|
|
if len(image.shape) == 3 and image.shape[2] == 3: # color image |
|
color = True |
|
elif (len(image.shape) == 2 or len(image.shape) == 3 and |
|
image.shape[2] == 1): # greyscale |
|
color = False |
|
else: |
|
raise Exception( |
|
"Image must have H x W x 3, H x W x 1 or H x W dimensions.") |
|
|
|
file.write("PF\n" if color else "Pf\n".encode()) |
|
file.write("%d %d\n".encode() % (image.shape[1], image.shape[0])) |
|
|
|
endian = image.dtype.byteorder |
|
|
|
if endian == "<" or endian == "=" and sys.byteorder == "little": |
|
scale = -scale |
|
|
|
file.write("%f\n".encode() % scale) |
|
|
|
image.tofile(file) |
|
|
|
|
|
def read_image(path): |
|
"""Read image and output RGB image (0-1). |
|
|
|
Args: |
|
path (str): path to file |
|
|
|
Returns: |
|
array: RGB image (0-1) |
|
""" |
|
img = cv2.imread(path) |
|
if img.ndim == 2: |
|
img = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR) |
|
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) / 255.0 |
|
return img |
|
|
|
|
|
def write_depth(path, depth, bits=1): |
|
"""Write depth map to pfm and png file. |
|
|
|
Args: |
|
path (str): filepath without extension |
|
depth (array): depth |
|
""" |
|
write_pfm(path + ".pfm", depth.astype(np.float32)) |
|
|
|
depth_min = depth.min() |
|
depth_max = depth.max() |
|
|
|
max_val = (2**(8 * bits)) - 1 |
|
|
|
if depth_max - depth_min > np.finfo("float").eps: |
|
out = max_val * (depth - depth_min) / (depth_max - depth_min) |
|
else: |
|
out = np.zeros(depth.shape, dtype=depth.type) |
|
|
|
if bits == 1: |
|
cv2.imwrite(path + ".png", out.astype("uint8")) |
|
elif bits == 2: |
|
cv2.imwrite(path + ".png", out.astype("uint16")) |
|
return path + '.pfm', path + ".png"
|
|
|