blanchon/Metric3D

1import matplotlib.pyplot as plt2import os, cv23import numpy as np4from mono.utils.transform import gray_to_colormap5import shutil6import glob7from mono.utils.running import main_process8import torch9from html4vision import Col, imagetable10 11def save_raw_imgs( 12    pred: torch.tensor,  13    rgb: torch.tensor, 14    filename: str, 15    save_dir: str,16    scale: float=200.0, 17    target: torch.tensor=None,18    ):19    """20    Save raw GT, predictions, RGB in the same file.21    """22    cv2.imwrite(os.path.join(save_dir, filename[:-4]+'_rgb.jpg'), rgb)23    cv2.imwrite(os.path.join(save_dir, filename[:-4]+'_d.png'), (pred*scale).astype(np.uint16))24    if target is not None:25        cv2.imwrite(os.path.join(save_dir, filename[:-4]+'_gt.png'), (target*scale).astype(np.uint16))26    27 28def save_val_imgs(29    iter: int, 30    pred: torch.tensor, 31    target: torch.tensor,32    rgb: torch.tensor, 33    filename: str, 34    save_dir: str, 35    tb_logger=None36    ):37    """38    Save GT, predictions, RGB in the same file.39    """40    rgb, pred_scale, target_scale, pred_color, target_color = get_data_for_log(pred, target, rgb)41    rgb = rgb.transpose((1, 2, 0))42    cat_img = np.concatenate([rgb, pred_color, target_color], axis=0)43    plt.imsave(os.path.join(save_dir, filename[:-4]+'_merge.jpg'), cat_img)44 45    # save to tensorboard46    if tb_logger is not None:47        tb_logger.add_image(f'{filename[:-4]}_merge.jpg', cat_img.transpose((2, 0, 1)), iter)48 49def save_normal_val_imgs(50    iter: int, 51    pred: torch.tensor, 52    targ: torch.tensor, 53    rgb: torch.tensor, 54    filename: str, 55    save_dir: str, 56    tb_logger=None, 57    mask=None,58    ):59    """60    Save GT, predictions, RGB in the same file.61    """62    mean = np.array([123.675, 116.28, 103.53])[np.newaxis, np.newaxis, :]63    std= np.array([58.395, 57.12, 57.375])[np.newaxis, np.newaxis, :]64    pred = pred.squeeze()65    targ = targ.squeeze()66    rgb = rgb.squeeze()67 68    if pred.size(0) == 3:69        pred = pred.permute(1,2,0)70    if targ.size(0) == 3:71        targ = targ.permute(1,2,0)72    if rgb.size(0) == 3:73        rgb = rgb.permute(1,2,0)74 75    pred_color = vis_surface_normal(pred, mask)76    targ_color = vis_surface_normal(targ, mask)77    rgb_color = ((rgb.cpu().numpy() * std) + mean).astype(np.uint8)78 79    try:80        cat_img = np.concatenate([rgb_color, pred_color, targ_color], axis=0)81    except:82        pred_color = cv2.resize(pred_color, (rgb.shape[1], rgb.shape[0]))83        targ_color = cv2.resize(targ_color, (rgb.shape[1], rgb.shape[0]))84        cat_img = np.concatenate([rgb_color, pred_color, targ_color], axis=0)85 86    plt.imsave(os.path.join(save_dir, filename[:-4]+'_merge.jpg'), cat_img)87    # cv2.imwrite(os.path.join(save_dir, filename[:-4]+'.jpg'), pred_color)88    # save to tensorboard89    if tb_logger is not None:90        tb_logger.add_image(f'{filename[:-4]}_merge.jpg', cat_img.transpose((2, 0, 1)), iter)91 92def get_data_for_log(pred: torch.tensor, target: torch.tensor, rgb: torch.tensor):93    mean = np.array([123.675, 116.28, 103.53])[:, np.newaxis, np.newaxis]94    std= np.array([58.395, 57.12, 57.375])[:, np.newaxis, np.newaxis]95 96    pred = pred.squeeze().cpu().numpy()97    target = target.squeeze().cpu().numpy()98    rgb = rgb.squeeze().cpu().numpy()99 100    pred[pred<0] = 0101    target[target<0] = 0102    max_scale = max(pred.max(), target.max())103    pred_scale = (pred/max_scale * 10000).astype(np.uint16)104    target_scale = (target/max_scale * 10000).astype(np.uint16)105    pred_color = gray_to_colormap(pred)106    target_color = gray_to_colormap(target)107    pred_color = cv2.resize(pred_color, (rgb.shape[2], rgb.shape[1]))108    target_color = cv2.resize(target_color, (rgb.shape[2], rgb.shape[1]))109 110    rgb = ((rgb * std) + mean).astype(np.uint8)111    return rgb, pred_scale, target_scale, pred_color, target_color112 113 114def create_html(name2path, save_path='index.html', size=(256, 384)):115    # table description116    cols = []117    for k, v in name2path.items():118        col_i =  Col('img', k, v) # specify image content for column119        cols.append(col_i)120    # html table generation121    imagetable(cols, out_file=save_path, imsize=size)122 123def vis_surface_normal(normal: torch.tensor, mask: torch.tensor=None) -> np.array:124    """125    Visualize surface normal. Transfer surface normal value from [-1, 1] to [0, 255]126    Aargs:127        normal (torch.tensor, [h, w, 3]): surface normal128        mask (torch.tensor, [h, w]): valid masks129    """130    normal = normal.cpu().numpy().squeeze()131    n_img_L2 = np.sqrt(np.sum(normal ** 2, axis=2, keepdims=True))132    n_img_norm = normal / (n_img_L2 + 1e-8)133    normal_vis = n_img_norm * 127134    normal_vis += 128135    normal_vis = normal_vis.astype(np.uint8)136    if mask is not None:137        mask = mask.cpu().numpy().squeeze()138        normal_vis[~mask] = 0139    return normal_vis140 141