import torch import torch.nn.functional as F import logging import os import os.path as osp from mono.utils.avg_meter import MetricAverageMeter from mono.utils.visualization import save_val_imgs, create_html, save_raw_imgs, save_normal_val_imgs import cv2 from tqdm import tqdm import numpy as np from PIL import Image import matplotlib.pyplot as plt from mono.utils.unproj_pcd import reconstruct_pcd, save_point_cloud def to_cuda(data: dict): for k, v in data.items(): if isinstance(v, torch.Tensor): data[k] = v.cuda(non_blocking=True) if isinstance(v, list) and len(v)>=1 and isinstance(v[0], torch.Tensor): for i, l_i in enumerate(v): data[k][i] = l_i.cuda(non_blocking=True) return data def align_scale(pred: torch.tensor, target: torch.tensor): mask = target > 0 if torch.sum(mask) > 10: scale = torch.median(target[mask]) / (torch.median(pred[mask]) + 1e-8) else: scale = 1 pred_scaled = pred * scale return pred_scaled, scale def align_scale_shift(pred: torch.tensor, target: torch.tensor): mask = target > 0 target_mask = target[mask].cpu().numpy() pred_mask = pred[mask].cpu().numpy() if torch.sum(mask) > 10: scale, shift = np.polyfit(pred_mask, target_mask, deg=1) if scale < 0: scale = torch.median(target[mask]) / (torch.median(pred[mask]) + 1e-8) shift = 0 else: scale = 1 shift = 0 pred = pred * scale + shift return pred, scale def align_scale_shift_numpy(pred: np.array, target: np.array): mask = target > 0 target_mask = target[mask] pred_mask = pred[mask] if np.sum(mask) > 10: scale, shift = np.polyfit(pred_mask, target_mask, deg=1) if scale < 0: scale = np.median(target[mask]) / (np.median(pred[mask]) + 1e-8) shift = 0 else: scale = 1 shift = 0 pred = pred * scale + shift return pred, scale def build_camera_model(H : int, W : int, intrinsics : list) -> np.array: """ Encode the camera intrinsic parameters (focal length and principle point) to a 4-channel map. """ fx, fy, u0, v0 = intrinsics f = (fx + fy) / 2.0 # principle point location x_row = np.arange(0, W).astype(np.float32) x_row_center_norm = (x_row - u0) / W x_center = np.tile(x_row_center_norm, (H, 1)) # [H, W] y_col = np.arange(0, H).astype(np.float32) y_col_center_norm = (y_col - v0) / H y_center = np.tile(y_col_center_norm, (W, 1)).T # [H, W] # FoV fov_x = np.arctan(x_center / (f / W)) fov_y = np.arctan(y_center / (f / H)) cam_model = np.stack([x_center, y_center, fov_x, fov_y], axis=2) return cam_model def resize_for_input(image, output_shape, intrinsic, canonical_shape, to_canonical_ratio): """ Resize the input. 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 the label will be scaled with the resize factor. """ padding = [123.675, 116.28, 103.53] h, w, _ = image.shape resize_ratio_h = output_shape[0] / canonical_shape[0] resize_ratio_w = output_shape[1] / canonical_shape[1] to_scale_ratio = min(resize_ratio_h, resize_ratio_w) resize_ratio = to_canonical_ratio * to_scale_ratio reshape_h = int(resize_ratio * h) reshape_w = int(resize_ratio * w) pad_h = max(output_shape[0] - reshape_h, 0) pad_w = max(output_shape[1] - reshape_w, 0) pad_h_half = int(pad_h / 2) pad_w_half = int(pad_w / 2) # resize image = cv2.resize(image, dsize=(reshape_w, reshape_h), interpolation=cv2.INTER_LINEAR) # padding image = cv2.copyMakeBorder( image, pad_h_half, pad_h - pad_h_half, pad_w_half, pad_w - pad_w_half, cv2.BORDER_CONSTANT, value=padding) # Resize, adjust principle point intrinsic[2] = intrinsic[2] * to_scale_ratio intrinsic[3] = intrinsic[3] * to_scale_ratio cam_model = build_camera_model(reshape_h, reshape_w, intrinsic) cam_model = cv2.copyMakeBorder( cam_model, pad_h_half, pad_h - pad_h_half, pad_w_half, pad_w - pad_w_half, cv2.BORDER_CONSTANT, value=-1) pad=[pad_h_half, pad_h - pad_h_half, pad_w_half, pad_w - pad_w_half] label_scale_factor=1/to_scale_ratio return image, cam_model, pad, label_scale_factor def get_prediction( model: torch.nn.Module, input: torch.tensor, cam_model: torch.tensor, pad_info: torch.tensor, scale_info: torch.tensor, gt_depth: torch.tensor, normalize_scale: float, ori_shape: list=[], ): data = dict( input=input, cam_model=cam_model, ) pred_depth, confidence, output_dict = model.module.inference(data) pred_depth = pred_depth pred_depth = pred_depth.squeeze() pred_depth = pred_depth[pad_info[0] : pred_depth.shape[0] - pad_info[1], pad_info[2] : pred_depth.shape[1] - pad_info[3]] if gt_depth is not None: resize_shape = gt_depth.shape elif ori_shape != []: resize_shape = ori_shape else: resize_shape = pred_depth.shape pred_depth = torch.nn.functional.interpolate(pred_depth[None, None, :, :], resize_shape, mode='bilinear').squeeze() # to original size pred_depth = pred_depth * normalize_scale / scale_info if gt_depth is not None: pred_depth_scale, scale = align_scale(pred_depth, gt_depth) else: pred_depth_scale = None scale = None return pred_depth, pred_depth_scale, scale, output_dict def transform_test_data_scalecano(rgb, intrinsic, data_basic): """ Pre-process the input for forwarding. Employ `label scale canonical transformation.' Args: rgb: input rgb image. [H, W, 3] intrinsic: camera intrinsic parameter, [fx, fy, u0, v0] data_basic: predefined canonical space in configs. """ canonical_space = data_basic['canonical_space'] forward_size = data_basic.crop_size mean = torch.tensor([123.675, 116.28, 103.53]).float()[:, None, None] std = torch.tensor([58.395, 57.12, 57.375]).float()[:, None, None] # BGR to RGB rgb = cv2.cvtColor(rgb, cv2.COLOR_BGR2RGB) ori_h, ori_w, _ = rgb.shape ori_focal = (intrinsic[0] + intrinsic[1]) / 2 canonical_focal = canonical_space['focal_length'] cano_label_scale_ratio = canonical_focal / ori_focal canonical_intrinsic = [ intrinsic[0] * cano_label_scale_ratio, intrinsic[1] * cano_label_scale_ratio, intrinsic[2], intrinsic[3], ] # resize rgb, cam_model, pad, resize_label_scale_ratio = resize_for_input(rgb, forward_size, canonical_intrinsic, [ori_h, ori_w], 1.0) # label scale factor label_scale_factor = cano_label_scale_ratio * resize_label_scale_ratio rgb = torch.from_numpy(rgb.transpose((2, 0, 1))).float() rgb = torch.div((rgb - mean), std) rgb = rgb[None, :, :, :].cuda() cam_model = torch.from_numpy(cam_model.transpose((2, 0, 1))).float() cam_model = cam_model[None, :, :, :].cuda() cam_model_stacks = [ torch.nn.functional.interpolate(cam_model, size=(cam_model.shape[2]//i, cam_model.shape[3]//i), mode='bilinear', align_corners=False) for i in [2, 4, 8, 16, 32] ] return rgb, cam_model_stacks, pad, label_scale_factor def do_scalecano_test_with_custom_data( model: torch.nn.Module, cfg: dict, test_data: list, logger: logging.RootLogger, is_distributed: bool = True, local_rank: int = 0, ): show_dir = cfg.show_dir save_interval = 1 save_imgs_dir = show_dir + '/vis' os.makedirs(save_imgs_dir, exist_ok=True) save_pcd_dir = show_dir + '/pcd' os.makedirs(save_pcd_dir, exist_ok=True) normalize_scale = cfg.data_basic.depth_range[1] dam = MetricAverageMeter(['abs_rel', 'rmse', 'silog', 'delta1', 'delta2', 'delta3']) dam_median = MetricAverageMeter(['abs_rel', 'rmse', 'silog', 'delta1', 'delta2', 'delta3']) dam_global = MetricAverageMeter(['abs_rel', 'rmse', 'silog', 'delta1', 'delta2', 'delta3']) for i, an in tqdm(enumerate(test_data)): #for i, an in enumerate(test_data): print(an['rgb']) rgb_origin = cv2.imread(an['rgb'])[:, :, ::-1].copy() if an['depth'] is not None: gt_depth = cv2.imread(an['depth'], -1) gt_depth_scale = an['depth_scale'] gt_depth = gt_depth / gt_depth_scale gt_depth_flag = True else: gt_depth = None gt_depth_flag = False intrinsic = an['intrinsic'] if intrinsic is None: intrinsic = [1000.0, 1000.0, rgb_origin.shape[1]/2, rgb_origin.shape[0]/2] # intrinsic = [542.0, 542.0, 963.706, 760.199] print(intrinsic) rgb_input, cam_models_stacks, pad, label_scale_factor = transform_test_data_scalecano(rgb_origin, intrinsic, cfg.data_basic) pred_depth, pred_depth_scale, scale, output = get_prediction( model = model, input = rgb_input, cam_model = cam_models_stacks, pad_info = pad, scale_info = label_scale_factor, gt_depth = None, normalize_scale = normalize_scale, ori_shape=[rgb_origin.shape[0], rgb_origin.shape[1]], ) pred_depth = (pred_depth > 0) * (pred_depth < 300) * pred_depth if gt_depth_flag: pred_depth = torch.nn.functional.interpolate(pred_depth[None, None, :, :], (gt_depth.shape[0], gt_depth.shape[1]), mode='bilinear').squeeze() # to original size gt_depth = torch.from_numpy(gt_depth).cuda() pred_depth_median = pred_depth * gt_depth[gt_depth != 0].median() / pred_depth[gt_depth != 0].median() pred_global, _ = align_scale_shift(pred_depth, gt_depth) mask = (gt_depth > 1e-8) dam.update_metrics_gpu(pred_depth, gt_depth, mask, is_distributed) dam_median.update_metrics_gpu(pred_depth_median, gt_depth, mask, is_distributed) dam_global.update_metrics_gpu(pred_global, gt_depth, mask, is_distributed) print(gt_depth[gt_depth != 0].median() / pred_depth[gt_depth != 0].median(), ) if i % save_interval == 0: os.makedirs(osp.join(save_imgs_dir, an['folder']), exist_ok=True) rgb_torch = torch.from_numpy(rgb_origin).to(pred_depth.device).permute(2, 0, 1) mean = torch.tensor([123.675, 116.28, 103.53]).float()[:, None, None].to(rgb_torch.device) std = torch.tensor([58.395, 57.12, 57.375]).float()[:, None, None].to(rgb_torch.device) rgb_torch = torch.div((rgb_torch - mean), std) save_val_imgs( i, pred_depth, gt_depth if gt_depth is not None else torch.ones_like(pred_depth, device=pred_depth.device), rgb_torch, osp.join(an['folder'], an['filename']), save_imgs_dir, ) #save_raw_imgs(pred_depth.detach().cpu().numpy(), rgb_torch, osp.join(an['folder'], an['filename']), save_imgs_dir, 1000.0) # pcd pred_depth = pred_depth.detach().cpu().numpy() #pcd = reconstruct_pcd(pred_depth, intrinsic[0], intrinsic[1], intrinsic[2], intrinsic[3]) #os.makedirs(osp.join(save_pcd_dir, an['folder']), exist_ok=True) #save_point_cloud(pcd.reshape((-1, 3)), rgb_origin.reshape(-1, 3), osp.join(save_pcd_dir, an['folder'], an['filename'][:-4]+'.ply')) if an['intrinsic'] == None: #for r in [0.9, 1.0, 1.1]: for r in [1.0]: #for f in [600, 800, 1000, 1250, 1500]: for f in [1000]: pcd = reconstruct_pcd(pred_depth, f * r, f * (2-r), intrinsic[2], intrinsic[3]) fstr = '_fx_' + str(int(f * r)) + '_fy_' + str(int(f * (2-r))) os.makedirs(osp.join(save_pcd_dir, an['folder']), exist_ok=True) save_point_cloud(pcd.reshape((-1, 3)), rgb_origin.reshape(-1, 3), osp.join(save_pcd_dir, an['folder'], an['filename'][:-4] + fstr +'.ply')) if "normal_out_list" in output.keys(): normal_out_list = output['normal_out_list'] pred_normal = normal_out_list[0][:, :3, :, :] # (B, 3, H, W) H, W = pred_normal.shape[2:] pred_normal = pred_normal[:, :, pad[0]:H-pad[1], pad[2]:W-pad[3]] gt_normal = None #if gt_normal_flag: if False: pred_normal = torch.nn.functional.interpolate(pred_normal, size=gt_normal.shape[2:], mode='bilinear', align_corners=True) gt_normal = cv2.imread(norm_path) gt_normal = cv2.cvtColor(gt_normal, cv2.COLOR_BGR2RGB) gt_normal = np.array(gt_normal).astype(np.uint8) gt_normal = ((gt_normal.astype(np.float32) / 255.0) * 2.0) - 1.0 norm_valid_mask = (np.linalg.norm(gt_normal, axis=2, keepdims=True) > 0.5) gt_normal = gt_normal * norm_valid_mask gt_normal_mask = ~torch.all(gt_normal == 0, dim=1, keepdim=True) dam.update_normal_metrics_gpu(pred_normal, gt_normal, gt_normal_mask, cfg.distributed)# save valiad normal if i % save_interval == 0: save_normal_val_imgs(iter, pred_normal, gt_normal if gt_normal is not None else torch.ones_like(pred_normal, device=pred_normal.device), rgb_torch, # data['input'], osp.join(an['folder'], 'normal_'+an['filename']), save_imgs_dir, ) #if gt_depth_flag: if False: eval_error = dam.get_metrics() print('w/o match :', eval_error) eval_error_median = dam_median.get_metrics() print('median match :', eval_error_median) eval_error_global = dam_global.get_metrics() print('global match :', eval_error_global) else: print('missing gt_depth, only save visualizations...')