blanchon/Metric3D
0
1import torch2import torch.nn.functional as F3import logging4import os5import os.path as osp6from mono.utils.avg_meter import MetricAverageMeter7from mono.utils.visualization import save_val_imgs, create_html, save_raw_imgs, save_normal_val_imgs8import cv29from tqdm import tqdm10import numpy as np11from PIL import Image12import matplotlib.pyplot as plt13 14from mono.utils.unproj_pcd import reconstruct_pcd, save_point_cloud15 16def to_cuda(data: dict):17 for k, v in data.items():18 if isinstance(v, torch.Tensor):19 data[k] = v.cuda(non_blocking=True)20 if isinstance(v, list) and len(v)>=1 and isinstance(v[0], torch.Tensor):21 for i, l_i in enumerate(v):22 data[k][i] = l_i.cuda(non_blocking=True)23 return data24 25def align_scale(pred: torch.tensor, target: torch.tensor):26 mask = target > 027 if torch.sum(mask) > 10:28 scale = torch.median(target[mask]) / (torch.median(pred[mask]) + 1e-8)29 else:30 scale = 131 pred_scaled = pred * scale32 return pred_scaled, scale33 34def align_scale_shift(pred: torch.tensor, target: torch.tensor):35 mask = target > 036 target_mask = target[mask].cpu().numpy()37 pred_mask = pred[mask].cpu().numpy()38 if torch.sum(mask) > 10:39 scale, shift = np.polyfit(pred_mask, target_mask, deg=1)40 if scale < 0:41 scale = torch.median(target[mask]) / (torch.median(pred[mask]) + 1e-8)42 shift = 043 else:44 scale = 145 shift = 046 pred = pred * scale + shift47 return pred, scale48 49def align_scale_shift_numpy(pred: np.array, target: np.array):50 mask = target > 051 target_mask = target[mask]52 pred_mask = pred[mask]53 if np.sum(mask) > 10:54 scale, shift = np.polyfit(pred_mask, target_mask, deg=1)55 if scale < 0:56 scale = np.median(target[mask]) / (np.median(pred[mask]) + 1e-8)57 shift = 058 else:59 scale = 160 shift = 061 pred = pred * scale + shift62 return pred, scale63 64 65def build_camera_model(H : int, W : int, intrinsics : list) -> np.array:66 """67 Encode the camera intrinsic parameters (focal length and principle point) to a 4-channel map. 68 """69 fx, fy, u0, v0 = intrinsics70 f = (fx + fy) / 2.071 # principle point location72 x_row = np.arange(0, W).astype(np.float32)73 x_row_center_norm = (x_row - u0) / W74 x_center = np.tile(x_row_center_norm, (H, 1)) # [H, W]75 76 y_col = np.arange(0, H).astype(np.float32) 77 y_col_center_norm = (y_col - v0) / H78 y_center = np.tile(y_col_center_norm, (W, 1)).T # [H, W]79 80 # FoV81 fov_x = np.arctan(x_center / (f / W))82 fov_y = np.arctan(y_center / (f / H))83 84 cam_model = np.stack([x_center, y_center, fov_x, fov_y], axis=2)85 return cam_model86 87def resize_for_input(image, output_shape, intrinsic, canonical_shape, to_canonical_ratio):88 """89 Resize the input.90 Resizing consists of two processed, i.e. 1) to the canonical space (adjust the camera model); 2) resize the image while the camera model holds. Thus the91 label will be scaled with the resize factor.92 """93 padding = [123.675, 116.28, 103.53]94 h, w, _ = image.shape95 resize_ratio_h = output_shape[0] / canonical_shape[0]96 resize_ratio_w = output_shape[1] / canonical_shape[1]97 to_scale_ratio = min(resize_ratio_h, resize_ratio_w)98 99 resize_ratio = to_canonical_ratio * to_scale_ratio100 101 reshape_h = int(resize_ratio * h)102 reshape_w = int(resize_ratio * w)103 104 pad_h = max(output_shape[0] - reshape_h, 0)105 pad_w = max(output_shape[1] - reshape_w, 0)106 pad_h_half = int(pad_h / 2)107 pad_w_half = int(pad_w / 2)108 109 # resize110 image = cv2.resize(image, dsize=(reshape_w, reshape_h), interpolation=cv2.INTER_LINEAR)111 # padding112 image = cv2.copyMakeBorder(113 image, 114 pad_h_half, 115 pad_h - pad_h_half, 116 pad_w_half, 117 pad_w - pad_w_half, 118 cv2.BORDER_CONSTANT, 119 value=padding)120 121 # Resize, adjust principle point122 intrinsic[2] = intrinsic[2] * to_scale_ratio123 intrinsic[3] = intrinsic[3] * to_scale_ratio124 125 cam_model = build_camera_model(reshape_h, reshape_w, intrinsic)126 cam_model = cv2.copyMakeBorder(127 cam_model, 128 pad_h_half, 129 pad_h - pad_h_half, 130 pad_w_half, 131 pad_w - pad_w_half, 132 cv2.BORDER_CONSTANT, 133 value=-1)134 135 pad=[pad_h_half, pad_h - pad_h_half, pad_w_half, pad_w - pad_w_half]136 label_scale_factor=1/to_scale_ratio137 return image, cam_model, pad, label_scale_factor138 139 140def get_prediction(141 model: torch.nn.Module,142 input: torch.tensor,143 cam_model: torch.tensor,144 pad_info: torch.tensor,145 scale_info: torch.tensor,146 gt_depth: torch.tensor,147 normalize_scale: float,148 ori_shape: list=[],149):150 151 data = dict(152 input=input,153 cam_model=cam_model,154 )155 #pred_depth, confidence, output_dict = model.module.inference(data)156 pred_depth, confidence, output_dict = model.inference(data)157 pred_depth = pred_depth.squeeze()158 pred_depth = pred_depth[pad_info[0] : pred_depth.shape[0] - pad_info[1], pad_info[2] : pred_depth.shape[1] - pad_info[3]]159 confidence = confidence.squeeze()160 confidence = confidence[pad_info[0] : confidence.shape[0] - pad_info[1], pad_info[2] : confidence.shape[1] - pad_info[3]]161 if gt_depth is not None:162 resize_shape = gt_depth.shape163 elif ori_shape != []:164 resize_shape = ori_shape165 else:166 resize_shape = pred_depth.shape167 168 pred_depth = torch.nn.functional.interpolate(pred_depth[None, None, :, :], resize_shape, mode='bilinear').squeeze() # to original size169 pred_depth = pred_depth * normalize_scale / scale_info170 if gt_depth is not None:171 pred_depth_scale, scale = align_scale(pred_depth, gt_depth)172 else:173 pred_depth_scale = None174 scale = None175 176 return pred_depth, pred_depth_scale, scale, output_dict, confidence177 178def transform_test_data_scalecano(rgb, intrinsic, data_basic):179 """180 Pre-process the input for forwarding. Employ `label scale canonical transformation.'181 Args:182 rgb: input rgb image. [H, W, 3]183 intrinsic: camera intrinsic parameter, [fx, fy, u0, v0]184 data_basic: predefined canonical space in configs.185 """186 canonical_space = data_basic['canonical_space']187 forward_size = data_basic.crop_size188 mean = torch.tensor([123.675, 116.28, 103.53]).float()[:, None, None]189 std = torch.tensor([58.395, 57.12, 57.375]).float()[:, None, None]190 191 # BGR to RGB192 rgb = cv2.cvtColor(rgb, cv2.COLOR_BGR2RGB)193 194 ori_h, ori_w, _ = rgb.shape195 ori_focal = (intrinsic[0] + intrinsic[1]) / 2196 canonical_focal = canonical_space['focal_length']197 198 cano_label_scale_ratio = canonical_focal / ori_focal199 200 canonical_intrinsic = [201 intrinsic[0] * cano_label_scale_ratio,202 intrinsic[1] * cano_label_scale_ratio,203 intrinsic[2],204 intrinsic[3],205 ]206 207 # resize208 rgb, cam_model, pad, resize_label_scale_ratio = resize_for_input(rgb, forward_size, canonical_intrinsic, [ori_h, ori_w], 1.0)209 210 # label scale factor211 label_scale_factor = cano_label_scale_ratio * resize_label_scale_ratio212 213 rgb = torch.from_numpy(rgb.transpose((2, 0, 1))).float()214 rgb = torch.div((rgb - mean), std)215 rgb = rgb[None, :, :, :].cuda()216 #rgb = rgb[None, :, :, :]217 218 cam_model = torch.from_numpy(cam_model.transpose((2, 0, 1))).float()219 cam_model = cam_model[None, :, :, :].cuda()220 #cam_model = cam_model[None, :, :, :]221 cam_model_stacks = [222 torch.nn.functional.interpolate(cam_model, size=(cam_model.shape[2]//i, cam_model.shape[3]//i), mode='bilinear', align_corners=False)223 for i in [2, 4, 8, 16, 32]224 ]225 return rgb, cam_model_stacks, pad, label_scale_factor226 227def do_scalecano_test_with_custom_data(228 model: torch.nn.Module,229 cfg: dict,230 test_data: list,231 logger: logging.RootLogger,232 is_distributed: bool = True,233 local_rank: int = 0,234):235 236 show_dir = cfg.show_dir237 save_interval = 1238 save_imgs_dir = show_dir + '/vis'239 os.makedirs(save_imgs_dir, exist_ok=True)240 save_pcd_dir = show_dir + '/pcd'241 os.makedirs(save_pcd_dir, exist_ok=True)242 243 normalize_scale = cfg.data_basic.depth_range[1]244 dam = MetricAverageMeter(['abs_rel', 'rmse', 'silog', 'delta1', 'delta2', 'delta3'])245 dam_median = MetricAverageMeter(['abs_rel', 'rmse', 'silog', 'delta1', 'delta2', 'delta3'])246 dam_global = MetricAverageMeter(['abs_rel', 'rmse', 'silog', 'delta1', 'delta2', 'delta3'])247 248 for i, an in tqdm(enumerate(test_data)):249 #for i, an in enumerate(test_data):250 print(an['rgb'])251 rgb_origin = cv2.imread(an['rgb'])[:, :, ::-1].copy()252 if an['depth'] is not None:253 gt_depth = cv2.imread(an['depth'], -1)254 gt_depth_scale = an['depth_scale']255 gt_depth = gt_depth / gt_depth_scale256 gt_depth_flag = True257 else:258 gt_depth = None259 gt_depth_flag = False260 intrinsic = an['intrinsic']261 if intrinsic is None:262 intrinsic = [1000.0, 1000.0, rgb_origin.shape[1]/2, rgb_origin.shape[0]/2]263 # intrinsic = [542.0, 542.0, 963.706, 760.199]264 print(intrinsic)265 rgb_input, cam_models_stacks, pad, label_scale_factor = transform_test_data_scalecano(rgb_origin, intrinsic, cfg.data_basic)266 267 pred_depth, pred_depth_scale, scale, output = get_prediction(268 model = model,269 input = rgb_input,270 cam_model = cam_models_stacks,271 pad_info = pad,272 scale_info = label_scale_factor,273 gt_depth = None,274 normalize_scale = normalize_scale,275 ori_shape=[rgb_origin.shape[0], rgb_origin.shape[1]],276 )277 278 pred_depth = (pred_depth > 0) * (pred_depth < 300) * pred_depth279 if gt_depth_flag:280 281 pred_depth = torch.nn.functional.interpolate(pred_depth[None, None, :, :], (gt_depth.shape[0], gt_depth.shape[1]), mode='bilinear').squeeze() # to original size282 283 #gt_depth = torch.from_numpy(gt_depth).cuda()284 gt_depth = torch.from_numpy(gt_depth)285 286 pred_depth_median = pred_depth * gt_depth[gt_depth != 0].median() / pred_depth[gt_depth != 0].median()287 pred_global, _ = align_scale_shift(pred_depth, gt_depth)288 289 mask = (gt_depth > 1e-8)290 dam.update_metrics_gpu(pred_depth, gt_depth, mask, is_distributed)291 dam_median.update_metrics_gpu(pred_depth_median, gt_depth, mask, is_distributed)292 dam_global.update_metrics_gpu(pred_global, gt_depth, mask, is_distributed)293 print(gt_depth[gt_depth != 0].median() / pred_depth[gt_depth != 0].median(), )294 295 if i % save_interval == 0:296 os.makedirs(osp.join(save_imgs_dir, an['folder']), exist_ok=True)297 rgb_torch = torch.from_numpy(rgb_origin).to(pred_depth.device).permute(2, 0, 1)298 mean = torch.tensor([123.675, 116.28, 103.53]).float()[:, None, None].to(rgb_torch.device)299 std = torch.tensor([58.395, 57.12, 57.375]).float()[:, None, None].to(rgb_torch.device)300 rgb_torch = torch.div((rgb_torch - mean), std)301 302 save_val_imgs(303 i,304 pred_depth,305 gt_depth if gt_depth is not None else torch.ones_like(pred_depth, device=pred_depth.device),306 rgb_torch,307 osp.join(an['folder'], an['filename']),308 save_imgs_dir,309 )310 #save_raw_imgs(pred_depth.detach().cpu().numpy(), rgb_torch, osp.join(an['folder'], an['filename']), save_imgs_dir, 1000.0)311 312 # pcd313 pred_depth = pred_depth.detach().cpu().numpy()314 #pcd = reconstruct_pcd(pred_depth, intrinsic[0], intrinsic[1], intrinsic[2], intrinsic[3])315 #os.makedirs(osp.join(save_pcd_dir, an['folder']), exist_ok=True)316 #save_point_cloud(pcd.reshape((-1, 3)), rgb_origin.reshape(-1, 3), osp.join(save_pcd_dir, an['folder'], an['filename'][:-4]+'.ply'))317 318 if an['intrinsic'] == None:319 #for r in [0.9, 1.0, 1.1]:320 for r in [1.0]:321 #for f in [600, 800, 1000, 1250, 1500]:322 for f in [1000]:323 pcd = reconstruct_pcd(pred_depth, f * r, f * (2-r), intrinsic[2], intrinsic[3])324 fstr = '_fx_' + str(int(f * r)) + '_fy_' + str(int(f * (2-r)))325 os.makedirs(osp.join(save_pcd_dir, an['folder']), exist_ok=True)326 save_point_cloud(pcd.reshape((-1, 3)), rgb_origin.reshape(-1, 3), osp.join(save_pcd_dir, an['folder'], an['filename'][:-4] + fstr +'.ply'))327 328 if "normal_out_list" in output.keys():329 330 normal_out_list = output['normal_out_list'] 331 pred_normal = normal_out_list[0][:, :3, :, :] # (B, 3, H, W)332 H, W = pred_normal.shape[2:]333 pred_normal = pred_normal[:, :, pad[0]:H-pad[1], pad[2]:W-pad[3]]334 335 gt_normal = None336 #if gt_normal_flag:337 if False:338 pred_normal = torch.nn.functional.interpolate(pred_normal, size=gt_normal.shape[2:], mode='bilinear', align_corners=True) 339 gt_normal = cv2.imread(norm_path)340 gt_normal = cv2.cvtColor(gt_normal, cv2.COLOR_BGR2RGB) 341 gt_normal = np.array(gt_normal).astype(np.uint8)342 gt_normal = ((gt_normal.astype(np.float32) / 255.0) * 2.0) - 1.0343 norm_valid_mask = (np.linalg.norm(gt_normal, axis=2, keepdims=True) > 0.5)344 gt_normal = gt_normal * norm_valid_mask 345 gt_normal_mask = ~torch.all(gt_normal == 0, dim=1, keepdim=True)346 dam.update_normal_metrics_gpu(pred_normal, gt_normal, gt_normal_mask, cfg.distributed)# save valiad normal347 348 if i % save_interval == 0:349 save_normal_val_imgs(iter, 350 pred_normal, 351 gt_normal if gt_normal is not None else torch.ones_like(pred_normal, device=pred_normal.device),352 rgb_torch, # data['input'], 353 osp.join(an['folder'], 'normal_'+an['filename']), 354 save_imgs_dir,355 )356 357 358 #if gt_depth_flag:359 if False:360 eval_error = dam.get_metrics()361 print('w/o match :', eval_error)362 363 eval_error_median = dam_median.get_metrics()364 print('median match :', eval_error_median)365 366 eval_error_global = dam_global.get_metrics()367 print('global match :', eval_error_global)368 else:369 print('missing gt_depth, only save visualizations...')370