training/mono/utils/visualization.py

import matplotlib.pyplot as plt
import os, cv2
import numpy as np
from mono.utils.transform import gray_to_colormap
import shutil
import glob
from mono.utils.running import main_process
import torch
from html4vision import Col, imagetable

def save_raw_imgs( 
    pred: torch.tensor,  
    rgb: torch.tensor, 
    filename: str, 
    save_dir: str,
    scale: float=1000.0, 
    target: torch.tensor=None,
    ):
    """
    Save raw GT, predictions, RGB in the same file.
    """
    cv2.imwrite(os.path.join(save_dir, filename[:-4]+'_rgb.jpg'), rgb)
    cv2.imwrite(os.path.join(save_dir, filename[:-4]+'_gt.png'), (pred*scale).astype(np.uint16))
    if target is not None:
        cv2.imwrite(os.path.join(save_dir, filename[:-4]+'_gt.png'), (target*scale).astype(np.uint16))
    
def save_normal_val_imgs(
    iter: int, 
    pred: torch.tensor, 
    #targ: torch.tensor, 
    #rgb: torch.tensor, 
    filename: str, 
    save_dir: str, 
    tb_logger=None, 
    mask=None,
    ):
    """
    Save GT, predictions, RGB in the same file.
    """
    mean = np.array([123.675, 116.28, 103.53])[np.newaxis, np.newaxis, :]
    std= np.array([58.395, 57.12, 57.375])[np.newaxis, np.newaxis, :]
    pred = pred.squeeze()
    
    # if pred.size(0) == 3:
    #     pred = pred.permute(1,2,0)
    # pred_color = vis_surface_normal(pred, mask)

    # #save one image only
    # plt.imsave(os.path.join(save_dir, filename[:-4]+'.jpg'), pred_color)

    targ = targ.squeeze()
    rgb = rgb.squeeze()

    if pred.size(0) == 3:
        pred = pred.permute(1,2,0)
    if targ.size(0) == 3:
        targ = targ.permute(1,2,0)
    if rgb.size(0) == 3:
        rgb = rgb.permute(1,2,0)

    pred_color = vis_surface_normal(pred, mask)
    targ_color = vis_surface_normal(targ, mask)
    rgb_color = ((rgb.cpu().numpy() * std) + mean).astype(np.uint8)

    try:
        cat_img = np.concatenate([rgb_color, pred_color, targ_color], axis=0)
    except:
        pred_color = cv2.resize(pred_color, (rgb.shape[1], rgb.shape[0]))
        targ_color = cv2.resize(targ_color, (rgb.shape[1], rgb.shape[0]))
        cat_img = np.concatenate([rgb_color, pred_color, targ_color], axis=0)

    plt.imsave(os.path.join(save_dir, filename[:-4]+'_merge.jpg'), cat_img)
    # cv2.imwrite(os.path.join(save_dir, filename[:-4]+'.jpg'), pred_color)
    # save to tensorboard
    if tb_logger is not None:
        tb_logger.add_image(f'{filename[:-4]}_merge.jpg', cat_img.transpose((2, 0, 1)), iter)

    


def save_val_imgs(
    iter: int, 
    pred: torch.tensor, 
    target: torch.tensor, 
    rgb: torch.tensor, 
    filename: str, 
    save_dir: str, 
    tb_logger=None
    ):
    """
    Save GT, predictions, RGB in the same file.
    """
    rgb, pred_scale, target_scale, pred_color, target_color, max_scale = get_data_for_log(pred, target, rgb)
    rgb = rgb.transpose((1, 2, 0))
    # plt.imsave(os.path.join(save_dir, filename[:-4]+'_rgb.jpg'), rgb)
    # plt.imsave(os.path.join(save_dir, filename[:-4]+'_pred.png'), pred_scale, cmap='rainbow')
    # plt.imsave(os.path.join(save_dir, filename[:-4]+'_gt.png'), target_scale, cmap='rainbow')
    cat_img = np.concatenate([rgb, pred_color, target_color], axis=0)
    plt.imsave(os.path.join(save_dir, filename[:-4]+'_merge.jpg'), cat_img)

    # save to tensorboard
    if tb_logger is not None:
        # tb_logger.add_image(f'{filename[:-4]}_rgb.jpg', rgb, iter)
        # tb_logger.add_image(f'{filename[:-4]}_pred.jpg', gray_to_colormap(pred_scale).transpose((2, 0, 1)), iter)
        # tb_logger.add_image(f'{filename[:-4]}_gt.jpg', gray_to_colormap(target_scale).transpose((2, 0, 1)), iter)
        tb_logger.add_image(f'{filename[:-4]}_merge.jpg', cat_img.transpose((2, 0, 1)), iter)
    return max_scale

def get_data_for_log(pred: torch.tensor, target: torch.tensor, rgb: torch.tensor):
    mean = np.array([123.675, 116.28, 103.53])[:, np.newaxis, np.newaxis]
    std= np.array([58.395, 57.12, 57.375])[:, np.newaxis, np.newaxis]

    pred = pred.squeeze().cpu().numpy()
    target = target.squeeze().cpu().numpy()
    rgb = rgb.squeeze().cpu().numpy()

    pred[pred<0] = 0
    target[target<0] = 0
    #max_scale = max(pred.max(), target.max())
    max_scale = min(2.0 * target.max(), pred.max())
    pred[pred > max_scale] = max_scale

    pred_scale = (pred/max_scale * 10000).astype(np.uint16)
    target_scale = (target/max_scale * 10000).astype(np.uint16)
    pred_color = gray_to_colormap(pred, max_value=max_scale)
    target_color = gray_to_colormap(target, max_value=max_scale)
    
    dilate = True
    if dilate == True:
        k=np.ones((3,3),np.uint8)
        target_color=cv2.dilate(target_color,k,iterations=1)

    pred_color = cv2.resize(pred_color, (rgb.shape[2], rgb.shape[1]))
    target_color = cv2.resize(target_color, (rgb.shape[2], rgb.shape[1]))

    rgb = ((rgb * std) + mean).astype(np.uint8)
    return rgb, pred_scale, target_scale, pred_color, target_color, max_scale


def create_html(name2path, save_path='index.html', size=(256, 384)):
    # table description
    cols = []
    for k, v in name2path.items():
        col_i =  Col('img', k, v) # specify image content for column
        cols.append(col_i)
    # html table generation
    imagetable(cols, out_file=save_path, imsize=size)


def visual_train_data(gt_depth, rgb, filename, wkdir, replace=False, pred=None):
    gt_depth = gt_depth.cpu().squeeze().numpy()
    rgb = rgb.cpu().squeeze().numpy()

    mean = np.array([123.675, 116.28, 103.53])[:, np.newaxis, np.newaxis]
    std= np.array([58.395, 57.12, 57.375])[:, np.newaxis, np.newaxis]
    mask = gt_depth > 0
    
    rgb = ((rgb * std) + mean).astype(np.uint8).transpose((1, 2, 0))
    gt_vis = gray_to_colormap(gt_depth)
    if replace:
        rgb[mask] = gt_vis[mask]

    if pred is not None:
        pred_depth = pred.detach().cpu().squeeze().numpy()
        pred_vis = gray_to_colormap(pred_depth)

    merge = np.concatenate([rgb, gt_vis, pred_vis], axis=0)
    
    save_path = os.path.join(wkdir, 'test_train', filename)
    os.makedirs(os.path.dirname(save_path), exist_ok=True)
    plt.imsave(save_path, merge)


def create_dir_for_validate_meta(work_dir, iter_id):
    curr_folders = glob.glob(work_dir + '/online_val/*0')
    curr_folders = [i for i in curr_folders if os.path.isdir(i)]
    if len(curr_folders) > 8:
        curr_folders.sort()
        del_foler = curr_folders.pop(0)
        print(del_foler)
        if main_process():
            # only rank==0 do it
            if os.path.exists(del_foler):   
                shutil.rmtree(del_foler)
            if os.path.exists(del_foler + '.html'):
                os.remove(del_foler + '.html')
        
    save_val_meta_data_dir = os.path.join(work_dir, 'online_val', '%08d'%iter_id)
    os.makedirs(save_val_meta_data_dir, exist_ok=True)
    return save_val_meta_data_dir


def vis_surface_normal(normal: torch.tensor, mask: torch.tensor=None) -> np.array:
    """
    Visualize surface normal. Transfer surface normal value from [-1, 1] to [0, 255]
    Aargs:
        normal (torch.tensor, [h, w, 3]): surface normal
        mask (torch.tensor, [h, w]): valid masks
    """
    normal = normal.cpu().numpy().squeeze()
    n_img_L2 = np.sqrt(np.sum(normal ** 2, axis=2, keepdims=True))
    n_img_norm = normal / (n_img_L2 + 1e-8)
    normal_vis = n_img_norm * 127
    normal_vis += 128
    normal_vis = normal_vis.astype(np.uint8)
    if mask is not None:
        mask = mask.cpu().numpy().squeeze()
        normal_vis[~mask] = 0
    return normal_vis