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