blanchon/Metric3D
0
1import collections2import cv23import math4import numpy as np5import numbers6import random7import torch8 9import matplotlib10import matplotlib.cm11 12 13"""14Provides a set of Pytorch transforms that use OpenCV instead of PIL (Pytorch default)15for image manipulation.16"""17 18class Compose(object):19 # Composes transforms: transforms.Compose([transforms.RandScale([0.5, 2.0]), transforms.ToTensor()])20 def __init__(self, transforms):21 self.transforms = transforms22 23 def __call__(self, images, labels, intrinsics, cam_models=None, other_labels=None, transform_paras=None):24 for t in self.transforms:25 images, labels, intrinsics, cam_models, other_labels, transform_paras = t(images, labels, intrinsics, cam_models, other_labels, transform_paras)26 return images, labels, intrinsics, cam_models, other_labels, transform_paras27 28 29class ToTensor(object):30 # Converts numpy.ndarray (H x W x C) to a torch.FloatTensor of shape (C x H x W).31 def __init__(self, **kwargs):32 return33 def __call__(self, images, labels, intrinsics, cam_models=None, other_labels=None, transform_paras=None):34 if not isinstance(images, list) or not isinstance(labels, list) or not isinstance(intrinsics, list):35 raise (RuntimeError("transform.ToTensor() only handle inputs/labels/intrinsics lists."))36 if len(images) != len(intrinsics):37 raise (RuntimeError("Numbers of images and intrinsics are not matched."))38 if not isinstance(images[0], np.ndarray) or not isinstance(labels[0], np.ndarray):39 raise (RuntimeError("transform.ToTensor() only handle np.ndarray for the input and label."40 "[eg: data readed by cv2.imread()].\n"))41 if not isinstance(intrinsics[0], list):42 raise (RuntimeError("transform.ToTensor() only handle list for the camera intrinsics"))43 44 if len(images[0].shape) > 3 or len(images[0].shape) < 2:45 raise (RuntimeError("transform.ToTensor() only handle image(np.ndarray) with 3 dims or 2 dims.\n"))46 if len(labels[0].shape) > 3 or len(labels[0].shape) < 2:47 raise (RuntimeError("transform.ToTensor() only handle label(np.ndarray) with 3 dims or 2 dims.\n"))48 49 if len(intrinsics[0]) >4 or len(intrinsics[0]) < 3:50 raise (RuntimeError("transform.ToTensor() only handle intrinsic(list) with 3 sizes or 4 sizes.\n"))51 52 for i, img in enumerate(images):53 if len(img.shape) == 2:54 img = np.expand_dims(img, axis=2)55 images[i] = torch.from_numpy(img.transpose((2, 0, 1))).float()56 for i, lab in enumerate(labels):57 if len(lab.shape) == 2:58 lab = np.expand_dims(lab, axis=0)59 labels[i] = torch.from_numpy(lab).float()60 for i, intrinsic in enumerate(intrinsics):61 if len(intrinsic) == 3:62 intrinsic = [intrinsic[0],] + intrinsic63 intrinsics[i] = torch.tensor(intrinsic, dtype=torch.float)64 if cam_models is not None:65 for i, cam_model in enumerate(cam_models):66 cam_models[i] = torch.from_numpy(cam_model.transpose((2, 0, 1))).float() if cam_model is not None else None67 if other_labels is not None:68 for i, lab in enumerate(other_labels):69 if len(lab.shape) == 2:70 lab = np.expand_dims(lab, axis=0)71 other_labels[i] = torch.from_numpy(lab).float()72 return images, labels, intrinsics, cam_models, other_labels, transform_paras73 74 75class Normalize(object):76 # Normalize tensor with mean and standard deviation along channel: channel = (channel - mean) / std77 def __init__(self, mean, std=None, **kwargs):78 if std is None:79 assert len(mean) > 080 else:81 assert len(mean) == len(std)82 self.mean = torch.tensor(mean).float()[:, None, None]83 self.std = torch.tensor(std).float()[:, None, None] if std is not None \84 else torch.tensor([1.0, 1.0, 1.0]).float()[:, None, None]85 86 def __call__(self, images, labels, intrinsics, cam_models=None, other_labels=None, transform_paras=None):87 # if self.std is None:88 # # for t, m in zip(image, self.mean):89 # # t.sub(m)90 # image = image - self.mean91 # if ref_images is not None:92 # for i, ref_i in enumerate(ref_images):93 # ref_images[i] = ref_i - self.mean94 # else:95 # # for t, m, s in zip(image, self.mean, self.std):96 # # t.sub(m).div(s)97 # image = (image - self.mean) / self.std98 # if ref_images is not None:99 # for i, ref_i in enumerate(ref_images):100 # ref_images[i] = (ref_i - self.mean) / self.std101 for i, img in enumerate(images):102 img = torch.div((img - self.mean), self.std)103 images[i] = img104 return images, labels, intrinsics, cam_models, other_labels, transform_paras105 106 107class LableScaleCanonical(object):108 """109 To solve the ambiguity observation for the mono branch, i.e. different focal length (object size) with the same depth, cameras are110 mapped to a canonical space. To mimic this, we set the focal length to a canonical one and scale the depth value. NOTE: resize the image based on the ratio can also solve111 Args:112 images: list of RGB images.113 labels: list of depth/disparity labels.114 other labels: other labels, such as instance segmentations, semantic segmentations...115 """116 def __init__(self, **kwargs):117 self.canonical_focal = kwargs['focal_length']118 119 def _get_scale_ratio(self, intrinsic):120 target_focal_x = intrinsic[0]121 label_scale_ratio = self.canonical_focal / target_focal_x122 pose_scale_ratio = 1.0123 return label_scale_ratio, pose_scale_ratio124 125 def __call__(self, images, labels, intrinsics, cam_models=None, other_labels=None, transform_paras=None):126 assert len(images[0].shape) == 3 and len(labels[0].shape) == 2127 assert labels[0].dtype == np.float32128 129 label_scale_ratio = None130 pose_scale_ratio = None131 132 for i in range(len(intrinsics)):133 img_i = images[i]134 label_i = labels[i] if i < len(labels) else None135 intrinsic_i = intrinsics[i].copy()136 cam_model_i = cam_models[i] if cam_models is not None and i < len(cam_models) else None137 138 label_scale_ratio, pose_scale_ratio = self._get_scale_ratio(intrinsic_i)139 140 # adjust the focal length, map the current camera to the canonical space141 intrinsics[i] = [intrinsic_i[0] * label_scale_ratio, intrinsic_i[1] * label_scale_ratio, intrinsic_i[2], intrinsic_i[3]]142 143 # scale the label to the canonical space144 if label_i is not None:145 labels[i] = label_i * label_scale_ratio146 147 if cam_model_i is not None:148 # As the focal length is adjusted (canonical focal length), the camera model should be re-built149 ori_h, ori_w, _ = img_i.shape150 cam_models[i] = build_camera_model(ori_h, ori_w, intrinsics[i])151 152 153 if transform_paras is not None:154 transform_paras.update(label_scale_factor=label_scale_ratio, focal_scale_factor=label_scale_ratio)155 156 return images, labels, intrinsics, cam_models, other_labels, transform_paras157 158 159class ResizeKeepRatio(object):160 """161 Resize and pad to a given size. Hold the aspect ratio.162 This resizing assumes that the camera model remains unchanged.163 Args:164 resize_size: predefined output size.165 """166 def __init__(self, resize_size, padding=None, ignore_label=-1, **kwargs):167 if isinstance(resize_size, int):168 self.resize_h = resize_size169 self.resize_w = resize_size170 elif isinstance(resize_size, collections.Iterable) and len(resize_size) == 2 \171 and isinstance(resize_size[0], int) and isinstance(resize_size[1], int) \172 and resize_size[0] > 0 and resize_size[1] > 0:173 self.resize_h = resize_size[0]174 self.resize_w = resize_size[1]175 else:176 raise (RuntimeError("crop size error.\n"))177 if padding is None:178 self.padding = padding179 elif isinstance(padding, list):180 if all(isinstance(i, numbers.Number) for i in padding):181 self.padding = padding182 else:183 raise (RuntimeError("padding in Crop() should be a number list\n"))184 if len(padding) != 3:185 raise (RuntimeError("padding channel is not equal with 3\n"))186 else:187 raise (RuntimeError("padding in Crop() should be a number list\n"))188 if isinstance(ignore_label, int):189 self.ignore_label = ignore_label190 else:191 raise (RuntimeError("ignore_label should be an integer number\n"))192 # self.crop_size = kwargs['crop_size']193 self.canonical_focal = kwargs['focal_length']194 195 def main_data_transform(self, image, label, intrinsic, cam_model, resize_ratio, padding, to_scale_ratio):196 """197 Resize data first and then do the padding.198 'label' will be scaled.199 """200 h, w, _ = image.shape201 reshape_h = int(resize_ratio * h)202 reshape_w = int(resize_ratio * w)203 204 pad_h, pad_w, pad_h_half, pad_w_half = padding205 206 # resize207 image = cv2.resize(image, dsize=(reshape_w, reshape_h), interpolation=cv2.INTER_LINEAR)208 # padding209 image = cv2.copyMakeBorder(210 image, 211 pad_h_half, 212 pad_h - pad_h_half, 213 pad_w_half, 214 pad_w - pad_w_half, 215 cv2.BORDER_CONSTANT, 216 value=self.padding)217 218 if label is not None:219 # label = cv2.resize(label, dsize=(reshape_w, reshape_h), interpolation=cv2.INTER_NEAREST)220 label = resize_depth_preserve(label, (reshape_h, reshape_w))221 label = cv2.copyMakeBorder(222 label, 223 pad_h_half, 224 pad_h - pad_h_half, 225 pad_w_half, 226 pad_w - pad_w_half, 227 cv2.BORDER_CONSTANT, 228 value=self.ignore_label)229 # scale the label230 label = label / to_scale_ratio231 232 # Resize, adjust principle point233 if intrinsic is not None:234 intrinsic[0] = intrinsic[0] * resize_ratio / to_scale_ratio235 intrinsic[1] = intrinsic[1] * resize_ratio / to_scale_ratio236 intrinsic[2] = intrinsic[2] * resize_ratio237 intrinsic[3] = intrinsic[3] * resize_ratio238 239 if cam_model is not None:240 #cam_model = cv2.resize(cam_model, dsize=(reshape_w, reshape_h), interpolation=cv2.INTER_LINEAR)241 cam_model = build_camera_model(reshape_h, reshape_w, intrinsic)242 cam_model = cv2.copyMakeBorder(243 cam_model, 244 pad_h_half, 245 pad_h - pad_h_half, 246 pad_w_half, 247 pad_w - pad_w_half, 248 cv2.BORDER_CONSTANT, 249 value=self.ignore_label)250 251 # Pad, adjust the principle point252 if intrinsic is not None:253 intrinsic[2] = intrinsic[2] + pad_w_half254 intrinsic[3] = intrinsic[3] + pad_h_half255 return image, label, intrinsic, cam_model256 257 def get_label_scale_factor(self, image, intrinsic, resize_ratio):258 ori_h, ori_w, _ = image.shape259 # crop_h, crop_w = self.crop_size260 ori_focal = intrinsic[0]261 262 to_canonical_ratio = self.canonical_focal / ori_focal263 to_scale_ratio = resize_ratio / to_canonical_ratio264 return to_scale_ratio265 266 def __call__(self, images, labels, intrinsics, cam_models=None, other_labels=None, transform_paras=None):267 target_h, target_w, _ = images[0].shape268 resize_ratio_h = self.resize_h / target_h269 resize_ratio_w = self.resize_w / target_w270 resize_ratio = min(resize_ratio_h, resize_ratio_w)271 reshape_h = int(resize_ratio * target_h)272 reshape_w = int(resize_ratio * target_w)273 pad_h = max(self.resize_h - reshape_h, 0)274 pad_w = max(self.resize_w - reshape_w, 0)275 pad_h_half = int(pad_h / 2)276 pad_w_half = int(pad_w / 2)277 278 pad_info = [pad_h, pad_w, pad_h_half, pad_w_half]279 to_scale_ratio = self.get_label_scale_factor(images[0], intrinsics[0], resize_ratio)280 281 for i in range(len(images)):282 img = images[i]283 label = labels[i] if i < len(labels) else None284 intrinsic = intrinsics[i] if i < len(intrinsics) else None285 cam_model = cam_models[i] if cam_models is not None and i < len(cam_models) else None286 img, label, intrinsic, cam_model = self.main_data_transform(287 img, label, intrinsic, cam_model, resize_ratio, pad_info, to_scale_ratio)288 images[i] = img289 if label is not None:290 labels[i] = label291 if intrinsic is not None:292 intrinsics[i] = intrinsic293 if cam_model is not None:294 cam_models[i] = cam_model295 296 if other_labels is not None:297 298 for i, other_lab in enumerate(other_labels):299 # resize300 other_lab = cv2.resize(other_lab, dsize=(reshape_w, reshape_h), interpolation=cv2.INTER_NEAREST)301 # pad302 other_labels[i] = cv2.copyMakeBorder(303 other_lab, 304 pad_h_half, 305 pad_h - pad_h_half, 306 pad_w_half, 307 pad_w - pad_w_half, 308 cv2.BORDER_CONSTANT, 309 value=self.ignore_label)310 311 pad = [pad_h_half, pad_h - pad_h_half, pad_w_half, pad_w - pad_w_half]312 if transform_paras is not None:313 pad_old = transform_paras['pad'] if 'pad' in transform_paras else [0,0,0,0]314 new_pad = [pad_old[0] + pad[0], pad_old[1] + pad[1], pad_old[2] + pad[2], pad_old[3] + pad[3]]315 transform_paras.update(dict(pad=new_pad))316 if 'label_scale_factor' in transform_paras:317 transform_paras['label_scale_factor'] = transform_paras['label_scale_factor'] * 1.0 / to_scale_ratio318 else:319 transform_paras.update(label_scale_factor=1.0/to_scale_ratio)320 return images, labels, intrinsics, cam_models, other_labels, transform_paras321 322 323class BGR2RGB(object):324 # Converts image from BGR order to RGB order, for model initialized from Pytorch325 def __init__(self, **kwargs):326 return327 def __call__(self, images, labels, intrinsics, cam_models=None,other_labels=None, transform_paras=None):328 for i, img in enumerate(images):329 images[i] = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)330 return images, labels, intrinsics, cam_models, other_labels, transform_paras331 332 333def resize_depth_preserve(depth, shape):334 """335 Resizes depth map preserving all valid depth pixels336 Multiple downsampled points can be assigned to the same pixel.337 338 Parameters339 ----------340 depth : np.array [h,w]341 Depth map342 shape : tuple (H,W)343 Output shape344 345 Returns346 -------347 depth : np.array [H,W,1]348 Resized depth map349 """350 # Store dimensions and reshapes to single column351 depth = np.squeeze(depth)352 h, w = depth.shape353 x = depth.reshape(-1)354 # Create coordinate grid355 uv = np.mgrid[:h, :w].transpose(1, 2, 0).reshape(-1, 2)356 # Filters valid points357 idx = x > 0358 crd, val = uv[idx], x[idx]359 # Downsamples coordinates360 crd[:, 0] = (crd[:, 0] * (shape[0] / h) + 0.5).astype(np.int32)361 crd[:, 1] = (crd[:, 1] * (shape[1] / w) + 0.5).astype(np.int32)362 # Filters points inside image363 idx = (crd[:, 0] < shape[0]) & (crd[:, 1] < shape[1])364 crd, val = crd[idx], val[idx]365 # Creates downsampled depth image and assigns points366 depth = np.zeros(shape)367 depth[crd[:, 0], crd[:, 1]] = val368 # Return resized depth map369 return depth370 371 372def build_camera_model(H : int, W : int, intrinsics : list) -> np.array:373 """374 Encode the camera intrinsic parameters (focal length and principle point) to a 4-channel map. 375 """376 fx, fy, u0, v0 = intrinsics377 f = (fx + fy) / 2.0378 # principle point location379 x_row = np.arange(0, W).astype(np.float32)380 x_row_center_norm = (x_row - u0) / W381 x_center = np.tile(x_row_center_norm, (H, 1)) # [H, W]382 383 y_col = np.arange(0, H).astype(np.float32) 384 y_col_center_norm = (y_col - v0) / H385 y_center = np.tile(y_col_center_norm, (W, 1)).T386 387 # FoV388 fov_x = np.arctan(x_center / (f / W))389 fov_y = np.arctan(y_center/ (f / H))390 391 cam_model = np.stack([x_center, y_center, fov_x, fov_y], axis=2)392 return cam_model393 394def gray_to_colormap(img, cmap='rainbow'):395 """396 Transfer gray map to matplotlib colormap397 """398 assert img.ndim == 2399 400 img[img<0] = 0401 mask_invalid = img < 1e-10402 #img = img / (img.max() + 1e-8)403 404 img_ = img.flatten()405 max_value = np.percentile(img_, q=98)406 img = img / (max_value + 1e-8)407 408 norm = matplotlib.colors.Normalize(vmin=0, vmax=1.1)409 cmap_m = matplotlib.cm.get_cmap(cmap)410 map = matplotlib.cm.ScalarMappable(norm=norm, cmap=cmap_m)411 colormap = (map.to_rgba(img)[:, :, :3] * 255).astype(np.uint8)412 colormap[mask_invalid] = 0413 return colormap