Spaces:
Runtime error
Runtime error
| """Utils for monoDepth. | |
| """ | |
| import sys | |
| import re | |
| import numpy as np | |
| import cv2 | |
| import torch | |
| import imageio | |
| def read_pfm(path): | |
| """Read pfm file. | |
| Args: | |
| path (str): path to file | |
| Returns: | |
| tuple: (data, scale) | |
| """ | |
| with open(path, "rb") as file: | |
| color = None | |
| width = None | |
| height = None | |
| scale = None | |
| endian = None | |
| header = file.readline().rstrip() | |
| if header.decode("ascii") == "PF": | |
| color = True | |
| elif header.decode("ascii") == "Pf": | |
| color = False | |
| else: | |
| raise Exception("Not a PFM file: " + path) | |
| dim_match = re.match(r"^(\d+)\s(\d+)\s$", file.readline().decode("ascii")) | |
| if dim_match: | |
| width, height = list(map(int, dim_match.groups())) | |
| else: | |
| raise Exception("Malformed PFM header.") | |
| scale = float(file.readline().decode("ascii").rstrip()) | |
| if scale < 0: | |
| # little-endian | |
| endian = "<" | |
| scale = -scale | |
| else: | |
| # big-endian | |
| endian = ">" | |
| data = np.fromfile(file, endian + "f") | |
| shape = (height, width, 3) if color else (height, width) | |
| data = np.reshape(data, shape) | |
| data = np.flipud(data) | |
| return data, scale | |
| 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 resize_image(img): | |
| """Resize image and make it fit for network. | |
| Args: | |
| img (array): image | |
| Returns: | |
| tensor: data ready for network | |
| """ | |
| height_orig = img.shape[0] | |
| width_orig = img.shape[1] | |
| unit_scale = 384. | |
| if width_orig > height_orig: | |
| scale = width_orig / unit_scale | |
| else: | |
| scale = height_orig / unit_scale | |
| height = (np.ceil(height_orig / scale / 32) * 32).astype(int) | |
| width = (np.ceil(width_orig / scale / 32) * 32).astype(int) | |
| img_resized = cv2.resize(img, (width, height), interpolation=cv2.INTER_AREA) | |
| img_resized = ( | |
| torch.from_numpy(np.transpose(img_resized, (2, 0, 1))).contiguous().float() | |
| ) | |
| img_resized = img_resized.unsqueeze(0) | |
| return img_resized | |
| def resize_depth(depth, width, height): | |
| """Resize depth map and bring to CPU (numpy). | |
| Args: | |
| depth (tensor): depth | |
| width (int): image width | |
| height (int): image height | |
| Returns: | |
| array: processed depth | |
| """ | |
| depth = torch.squeeze(depth[0, :, :, :]).to("cpu") | |
| depth = cv2.blur(depth.numpy(), (3, 3)) | |
| depth_resized = cv2.resize( | |
| depth, (width, height), interpolation=cv2.INTER_AREA | |
| ) | |
| return depth_resized | |
| 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 = 0 | |
| if bits == 1: | |
| cv2.imwrite(path + ".png", out.astype("uint8")) | |
| elif bits == 2: | |
| cv2.imwrite(path + ".png", out.astype("uint16")) | |
| return |