|
- import numpy as np
- import os, imageio, sys
-
- import torch
- from tqdm.auto import tqdm
- from utils.ray_utils import get_rays
- from utils.utils import *
- from utils.ray_utils import ndc_rays
- from utils.metrics import ConfusionMatrix, calculate_ap
-
- from sklearn.decomposition import PCA
-
-
- def OctreeRender_trilinear_fast(rays, tensorf, chunk=4096, N_samples=-1, ndc_ray=False, white_bg=True, is_train=False,
- fp16=False, device='cuda'):
- rgbs, alphas, depth_maps, weights, uncertainties = [], [], [], [], []
- N_rays_all = rays.shape[0]
- for chunk_idx in range(N_rays_all // chunk + int(N_rays_all % chunk > 0)):
- rays_chunk = rays[chunk_idx * chunk:(chunk_idx + 1) * chunk].to(device)
-
- with torch.cuda.amp.autocast(enabled=fp16):
- rgb_map, depth_map = tensorf(rays_chunk, is_train=is_train, white_bg=white_bg, ndc_ray=ndc_ray,
- N_samples=N_samples)
-
- rgbs.append(rgb_map)
- depth_maps.append(depth_map)
-
- return torch.cat(rgbs), None, torch.cat(depth_maps), None, None
-
-
- def OctreeRender_trilinear_fast_with_sem(rays, sem_tensorf, chunk=4096, N_samples=-1,
- ndc_ray=False, H=None, W=None, focal=None, near=None,
- white_bg=True, is_train=False, fp16=False, device='cuda',
- tps_function=None, use_rgbs=True, use_raw_semfeas=False):
- rgbs, alphas, sem_maps, raw_semfea_maps, depth_maps, weights, uncertainties = [], [], [], [], [], [], []
- N_rays_all = rays.shape[0]
- loss_3d = {}
- num_chunk = 0
- for chunk_idx in range(N_rays_all // chunk + int(N_rays_all % chunk > 0)):
- rays_chunk = rays[chunk_idx * chunk:(chunk_idx + 1) * chunk].to(device)
-
- with torch.cuda.amp.autocast(enabled=fp16):
- rgb_map, depth_map, sem_map, raw_insfea_map, loss_3d_chunk = sem_tensorf(rays_chunk,
- is_train=is_train,
- white_bg=white_bg,
- ndc_ray=ndc_ray,
- N_samples=N_samples,
- # fp16=fp16,
- tps_function=tps_function) # result of sem_tensorf.forward
- # use_raw_semfeas=use_raw_semfeas)
- num_chunk += 1
- for key in loss_3d_chunk.keys():
- if key not in loss_3d.keys():
- loss_3d[key] = 0.
- loss_3d[key] += loss_3d_chunk[key]
-
- rgbs.append(rgb_map)
- sem_maps.append(sem_map)
- raw_semfea_maps.append(raw_insfea_map)
- depth_maps.append(depth_map)
-
- for key in loss_3d.keys():
- loss_3d[key] /= num_chunk
-
- render_ret = [None, None, None, None, torch.cat(depth_maps), None, None, loss_3d]
-
- if use_rgbs:
- render_ret[0] = torch.cat(rgbs)
- if use_raw_semfeas:
- render_ret[3] = torch.cat(raw_semfea_maps)
- else:
- render_ret[2] = torch.cat(sem_maps)
-
- return tuple(render_ret)
-
- # if use_raw_semfeas:
- # return torch.cat(rgbs), None, None, torch.cat(raw_semfea_maps), torch.cat(depth_maps), None, None, loss_3d
- # else:
- # return torch.cat(rgbs), None, torch.cat(sem_maps), None, torch.cat(depth_maps), None, None, loss_3d
-
-
- def OctreeRender_trilinear_fast_mani(rays, sem_tensorf, N_samples=-1,
- ndc_ray=False, white_bg=True, is_train=False, fp16=False, device='cuda',
- pt_feas=None):
- loss_3d = {}
-
- with torch.cuda.amp.autocast(enabled=fp16):
- pt_feas, rgb_map, depth_map, sem_map, _, loss_3d_chunk = sem_tensorf.render_mani(rays,
- white_bg=white_bg,
- is_train=is_train,
- N_samples=N_samples,
- ndc_ray=ndc_ray,
- pt_feas=pt_feas)
- for key in loss_3d_chunk.keys():
- if key not in loss_3d.keys():
- loss_3d[key] = 0.
- loss_3d[key] += loss_3d_chunk[key]
-
- return pt_feas, rgb_map, sem_map, depth_map, loss_3d
-
-
- @torch.no_grad()
- def evaluation(test_dataset, tensorf, args, renderer, savePath=None, N_vis=5, prtx='', N_samples=-1,
- white_bg=False, ndc_ray=False, compute_extra_metrics=True, fp16=False, device='cuda', chunk_size=4096):
- PSNRs, rgb_maps, depth_maps = [], [], []
- ssims, l_alex, l_vgg = [], [], []
- os.makedirs(savePath, exist_ok=True)
- os.makedirs(savePath+"/rgbd", exist_ok=True)
- try:
- tqdm._instances.clear()
- except Exception:
- pass
-
- near_far = test_dataset.near_far
- img_eval_interval = 1 if N_vis < 0 else max(test_dataset.all_rays.shape[0] // N_vis, 1) # negative N_vis: all test_dataset
- idxs = list(range(0, test_dataset.all_rays.shape[0], img_eval_interval))
- for idx, samples in tqdm(enumerate(test_dataset.all_rays[0::img_eval_interval]), file=sys.stdout):
- W, H = test_dataset.img_wh
- rays = samples.view(-1,samples.shape[-1])
-
- rgb_map, _, depth_map, _, _ = renderer(rays, tensorf, chunk=chunk_size, N_samples=N_samples,
- ndc_ray=ndc_ray, white_bg=white_bg, fp16=fp16, device=device)
- rgb_map = rgb_map.clamp(0.0, 1.0)
- rgb_map, depth_map = rgb_map.reshape(H, W, 3).cpu(), depth_map.reshape(H, W).cpu()
- depth_map, _ = visualize_depth_numpy(depth_map.numpy(),near_far)
-
- if len(test_dataset.all_rgbs):
- gt_rgb = test_dataset.all_rgbs[idxs[idx]].view(H, W, 3)
- loss = torch.mean((rgb_map - gt_rgb) ** 2)
- PSNRs.append(-10.0 * np.log(loss.item()) / np.log(10.0))
-
- if compute_extra_metrics:
- ssim = rgb_ssim(rgb_map, gt_rgb, 1)
- l_a = rgb_lpips(gt_rgb.numpy(), rgb_map.numpy(), 'alex', tensorf.device)
- l_v = rgb_lpips(gt_rgb.numpy(), rgb_map.numpy(), 'vgg', tensorf.device)
- ssims.append(ssim)
- l_alex.append(l_a)
- l_vgg.append(l_v)
-
- rgb_map = (rgb_map.numpy() * 255).astype('uint8')
- # rgb_map = np.concatenate((rgb_map, depth_map), axis=1)
- rgb_maps.append(rgb_map)
- depth_maps.append(depth_map)
- if savePath is not None:
- imageio.imwrite(f'{savePath}/{prtx}{idx:03d}.png', rgb_map)
- rgb_map = np.concatenate((rgb_map, depth_map), axis=1)
- imageio.imwrite(f'{savePath}/rgbd/{prtx}{idx:03d}.png', rgb_map)
-
- imageio.mimwrite(f'{savePath}/{prtx}video.mp4', np.stack(rgb_maps), fps=30, quality=10)
- imageio.mimwrite(f'{savePath}/{prtx}depthvideo.mp4', np.stack(depth_maps), fps=30, quality=10)
-
- if PSNRs:
- psnr = np.mean(np.asarray(PSNRs))
- if compute_extra_metrics:
- ssim = np.mean(np.asarray(ssims))
- l_a = np.mean(np.asarray(l_alex))
- l_v = np.mean(np.asarray(l_vgg))
- np.savetxt(f'{savePath}/{prtx}mean.txt', np.asarray([psnr, ssim, l_a, l_v]))
- else:
- np.savetxt(f'{savePath}/{prtx}mean.txt', np.asarray([psnr]))
-
- return PSNRs
-
-
- @torch.no_grad()
- def evaluation_with_semfea(test_dataset, sem_tensorf, args, sem_renderer, savePath=None, N_vis=5, prtx='', N_samples=-1,
- white_bg=False, ndc_ray=False, compute_extra_metrics=True, fp16=False, device='cuda', chunk_size=4096,
- use_rgbs=True, use_raw_semfeas=False):
- PSNRs, rgb_maps, depth_maps = [], [], []
- ssims, l_alex, l_vgg = [], [], []
- os.makedirs(savePath, exist_ok=True)
- os.makedirs(savePath+"/rgbd", exist_ok=True)
- os.makedirs(savePath+"/semfea", exist_ok=True)
- try:
- tqdm._instances.clear()
- except Exception:
- pass
-
- near_far = test_dataset.near_far
- img_eval_interval = 1 if N_vis < 0 else max(test_dataset.all_rays.shape[0] // N_vis, 1) # negative N_vis: all test_dataset
- idxs = list(range(0, test_dataset.all_rays.shape[0], img_eval_interval))
- for idx, samples in tqdm(enumerate(test_dataset.all_rays[0::img_eval_interval]), file=sys.stdout):
- W, H = test_dataset.img_wh
- rays = samples.view(-1, samples.shape[-1])
-
- rgb_map, _, _, raw_semfea_map, depth_map, _, _, loss_3d = sem_renderer(rays, sem_tensorf,
- chunk=chunk_size, N_samples=N_samples,
- ndc_ray=ndc_ray, white_bg=white_bg,
- fp16=fp16, device=device,
- use_rgbs=use_rgbs, use_raw_semfeas=use_raw_semfeas)
- if use_rgbs:
- rgb_map = rgb_map.clamp(0.0, 1.0)
- rgb_map = rgb_map.reshape(H, W, 3).cpu()
- depth_map = depth_map.reshape(H, W).cpu()
- depth_map, _ = visualize_depth_numpy(depth_map.numpy(), near_far)
-
- if len(test_dataset.all_rgbs):
- gt_rgb = test_dataset.all_rgbs[idxs[idx]].view(H, W, 3)
- loss = torch.mean((rgb_map - gt_rgb) ** 2)
- PSNRs.append(-10.0 * np.log(loss.item()) / np.log(10.0))
-
- if compute_extra_metrics:
- ssim = rgb_ssim(rgb_map, gt_rgb, 1)
- l_a = rgb_lpips(gt_rgb.numpy(), rgb_map.numpy(), 'alex', sem_tensorf.device)
- l_v = rgb_lpips(gt_rgb.numpy(), rgb_map.numpy(), 'vgg', sem_tensorf.device)
- ssims.append(ssim)
- l_alex.append(l_a)
- l_vgg.append(l_v)
-
- rgb_map = (rgb_map.numpy() * 255).astype('uint8')
- # rgb_map = np.concatenate((rgb_map, depth_map), axis=1)
- rgb_maps.append(rgb_map)
- depth_maps.append(depth_map)
- if savePath is not None:
- imageio.imwrite(f'{savePath}/{prtx}{idx:03d}.png', rgb_map)
- rgb_map = np.concatenate((rgb_map, depth_map), axis=1)
- imageio.imwrite(f'{savePath}/rgbd/{prtx}{idx:03d}.png', rgb_map)
-
- if use_raw_semfeas:
- raw_semfea_map = raw_semfea_map.reshape(H*W, -1).cpu().numpy()
- pca = PCA(n_components=3)
- pca.fit(raw_semfea_map)
- vis_semfea_map = pca.transform(raw_semfea_map)
- vis_semfea_map = np.reshape(vis_semfea_map, (H, W, 3))
- vis_semfea_map = np.uint8(((vis_semfea_map + 1.) / 2.) * 255.)
- if savePath is not None:
- imageio.imwrite(f'{savePath}/semfea/{prtx}{idx:03d}.png', vis_semfea_map)
-
- if use_rgbs:
- imageio.mimwrite(f'{savePath}/{prtx}video.mp4', np.stack(rgb_maps), fps=30, quality=10)
- imageio.mimwrite(f'{savePath}/{prtx}depthvideo.mp4', np.stack(depth_maps), fps=30, quality=10)
-
- if PSNRs:
- psnr = np.mean(np.asarray(PSNRs))
- if compute_extra_metrics:
- ssim = np.mean(np.asarray(ssims))
- l_a = np.mean(np.asarray(l_alex))
- l_v = np.mean(np.asarray(l_vgg))
- np.savetxt(f'{savePath}/{prtx}mean.txt', np.asarray([psnr, ssim, l_a, l_v]))
- else:
- np.savetxt(f'{savePath}/{prtx}mean.txt', np.asarray([psnr]))
-
- return PSNRs
- else:
- return None
-
-
- @torch.no_grad()
- def evaluation_with_sem(test_dataset, sem_tensorf, args, sem_renderer, savePath=None, N_vis=5, prtx='', N_samples=-1,
- white_bg=False, ndc_ray=False, compute_extra_metrics=True, fp16=False, device='cuda', chunk_size=4096,
- use_rgbs=True):
- PSNRs, rgb_maps, depth_maps, sem_maps = [], [], [], []
- ssims, l_alex, l_vgg = [], [], []
- os.makedirs(savePath, exist_ok=True)
- os.makedirs(savePath + "/rgbd", exist_ok=True)
- os.makedirs(savePath + "/sem", exist_ok=True)
-
- try:
- tqdm._instances.clear()
- except Exception:
- pass
-
- near_far = test_dataset.near_far
- img_eval_interval = 1 if N_vis < 0 else max(test_dataset.all_rays.shape[0] // N_vis, 1)
- idxs = list(range(0, test_dataset.all_rays.shape[0], img_eval_interval))
- for idx, samples in tqdm(enumerate(test_dataset.all_rays[0::img_eval_interval]), file=sys.stdout):
- W, H = test_dataset.img_wh
- rays = samples.view(-1, samples.shape[-1])
-
- rgb_map, _, sem_map, semfea_map, depth_map, _, _, loss_3d = sem_renderer(rays, sem_tensorf, # sem_map: (chunk_size, num_sem_classes)
- chunk=chunk_size,
- N_samples=N_samples,
- ndc_ray=ndc_ray,
- white_bg=white_bg,
- fp16=fp16,
- device=device,
- use_raw_semfeas=False,
- use_rgbs=use_rgbs)
-
- rgb_map = rgb_map.clamp(0.0, 1.0) # (H*W, 3)
- rgb_map, depth_map = rgb_map.reshape(H, W, 3).cpu(), depth_map.reshape(H, W).cpu()
-
- # sem_map = sem_map.clamp(0.0, 1.0)
- logits_2_label = lambda x: torch.argmax(torch.nn.functional.softmax(x, dim=-1), dim=-1)
- sem_map = logits_2_label(sem_map)
- sem_map = sem_map.cpu().numpy()
- sem_map_label = sem_map.copy() # ?
- """not the same as replica () or []"""
- sem_map = test_dataset.label_color_map(sem_map) # ndarray ?
- if args.dataset_name == "replica":
- # used in replica ?
- ignore_label = -1
- sem_map[sem_map_label == ignore_label, :] = 0
-
- sem_map = sem_map.reshape(H, W, 3).astype('uint8')
-
- depth_map, _ = visualize_depth_numpy(depth_map.numpy(),near_far)
- if len(test_dataset.all_rgbs):
- gt_rgb = test_dataset.all_rgbs[idxs[idx]].view(H, W, 3)
- loss = torch.mean((rgb_map - gt_rgb) ** 2)
- PSNRs.append(-10.0 * np.log(loss.item()) / np.log(10.0))
-
- if compute_extra_metrics:
- ssim = rgb_ssim(rgb_map, gt_rgb, 1)
- l_a = rgb_lpips(gt_rgb.numpy(), rgb_map.numpy(), 'alex', sem_tensorf.device)
- l_v = rgb_lpips(gt_rgb.numpy(), rgb_map.numpy(), 'vgg', sem_tensorf.device)
- ssims.append(ssim)
- l_alex.append(l_a)
- l_vgg.append(l_v)
-
- rgb_map = (rgb_map.numpy() * 255).astype('uint8')
- # rgb_map = np.concatenate((rgb_map, depth_map), axis=1)
- rgb_maps.append(rgb_map)
- sem_maps.append(sem_map)
- depth_maps.append(depth_map)
- if savePath is not None:
- imageio.imwrite(f'{savePath}/{prtx}{idx:03d}.png', rgb_map)
- rgb_map = np.concatenate((rgb_map, depth_map), axis=1)
- imageio.imwrite(f'{savePath}/rgbd/{prtx}{idx:03d}.png', rgb_map)
- imageio.imwrite(f'{savePath}/sem/{prtx}{idx:03d}.png', sem_map)
-
- imageio.mimwrite(f'{savePath}/{prtx}video.mp4', np.stack(rgb_maps), fps=30, quality=10)
- imageio.mimwrite(f'{savePath}/{prtx}depthvideo.mp4', np.stack(depth_maps), fps=30, quality=10)
- imageio.mimwrite(f'{savePath}/{prtx}sem_video.mp4', np.stack(sem_maps), fps=10, quality=10)
-
- if PSNRs:
- psnr = np.mean(np.asarray(PSNRs))
- if compute_extra_metrics:
- ssim = np.mean(np.asarray(ssims))
- l_a = np.mean(np.asarray(l_alex))
- l_v = np.mean(np.asarray(l_vgg))
- np.savetxt(f'{savePath}/{prtx}mean.txt', np.asarray([psnr, ssim, l_a, l_v]))
- else:
- np.savetxt(f'{savePath}/{prtx}mean.txt', np.asarray([psnr]))
-
- return PSNRs
-
-
- @torch.no_grad()
- def evaluation_sem(test_dataset, sem_tensorf, args, sem_renderer, savePath=None, N_vis=5, prtx='', N_samples=-1,
- white_bg=False, ndc_ray=False, compute_extra_metrics=True, fp16=False, device='cuda', chunk_size=4096):
- sem_maps = []
- os.makedirs(savePath, exist_ok=True)
- os.makedirs(savePath + "/sem", exist_ok=True)
-
- try:
- tqdm._instances.clear()
- except Exception:
- pass
-
- near_far = test_dataset.near_far
- img_eval_interval = 1 if N_vis < 0 else max(test_dataset.all_rays.shape[0] // N_vis, 1)
- idxs = list(range(0, test_dataset.all_rays.shape[0], img_eval_interval))
- if args.dataset_name == 'replica_dmnerf':
- gt_label_maps = test_dataset.sem_samples["sem_remap"]
- elif args.dataset_name == 'scannet':
- gt_label_maps = test_dataset.sem_samples["sem_img"]
-
- all_ap = []
- metric_ious = []
-
- sem_label_maps = []
- val_cm_pix = ConfusionMatrix(num_classes=test_dataset.num_semantic_class)
-
- for idx, samples in tqdm(enumerate(test_dataset.all_rays[0::img_eval_interval]), file=sys.stdout):
- W, H = test_dataset.img_wh
- num_semantic_class = test_dataset.num_semantic_class
- rays = samples.view(-1, samples.shape[-1])
-
- rgb_map, _, sem_map, semfea_map, depth_map, _, _, loss_3d = sem_renderer(rays, sem_tensorf, # sem_map: (chunk_size, num_sem_classes)
- chunk=chunk_size,
- N_samples=N_samples,
- ndc_ray=ndc_ray,
- white_bg=white_bg,
- fp16=fp16,
- device=device,
- use_raw_semfeas=False,
- use_rgbs=False)
-
- # rgb_map = rgb_map.clamp(0.0, 1.0) # (H*W, 3)
- # rgb_map, depth_map = rgb_map.reshape(H, W, 3).cpu(), depth_map.reshape(H, W).cpu()
-
- # sem_map = sem_map.clamp(0.0, 1.0)
- sem_logit_map = sem_map.cpu().numpy()
- logits_2_label = lambda x: torch.argmax(torch.nn.functional.softmax(x, dim=-1), dim=-1)
- sem_map = logits_2_label(sem_map)
- sem_map = sem_map.cpu().numpy()
- sem_label_map = sem_map.copy().reshape(-1, 1)
- """not the same as replica () or []"""
- sem_map = test_dataset.label_color_map(sem_map) # ndarray ?
- if args.dataset_name == "replica":
- # used in replica ?
- ignore_label = -1
- sem_map[sem_label_map == ignore_label, :] = 0
-
- # sem_logit_map = sem_logit_map.reshape(H, W, num_semantic_class)
- # sem_label_map = sem_label_map.reshape(H, W).astype('uint8')
- # sem_map = sem_map.reshape(H, W, 3).astype('uint8')
-
- gt_label_map = gt_label_maps[idx * img_eval_interval].reshape(-1, 1)
-
- val_cm = ConfusionMatrix(num_classes=test_dataset.num_semantic_class)
- metric_iou = val_cm.add_batch(sem_label_map, gt_label_map, return_miou=True)
- '''mIoU'''
- metric_ious.append(metric_iou)
-
- '''mAP'''
- confusion_matrix = val_cm.confusion_matrix
-
- unique_pred_labels = np.unique(sem_label_map)
- unique_gt_labels = np.unique(gt_label_map)
- num_valid_labels = len(unique_gt_labels)
- siou = np.divide(np.diag(confusion_matrix), (np.sum(confusion_matrix, axis=1) + np.sum(confusion_matrix, axis=0)
- - np.diag(confusion_matrix) + 1e-6))
- iou_metrics = siou[unique_gt_labels]
-
- '''confidence values'''
- # prepare confidence values
- unique_pred_labels, unique_gt_labels = torch.from_numpy(unique_pred_labels), torch.from_numpy(unique_gt_labels)
- pred_label_map = torch.from_numpy(sem_label_map).reshape(H, W)
- pred_conf_mask = torch.from_numpy(sem_logit_map).reshape(H, W, num_semantic_class)
- conf_scores = torch.zeros_like(unique_gt_labels, dtype=torch.float32)
- for i, label in enumerate(unique_gt_labels):
- if label.item() in unique_pred_labels:
- index = torch.where(pred_label_map == label)
- ssm = pred_conf_mask[index[0], index[1]] # confidence value
- pred_obj_conf = torch.median(ssm).item() # median confidence value for one object
- conf_scores[i] = pred_obj_conf
-
- iou_metrics = torch.from_numpy(iou_metrics).to(device)
- conf_scores = conf_scores.to(device)
-
- thre_list = [0.5, 0.75]
- ap = calculate_ap(iou_metrics, num_valid_labels, thre_list, device, confidence=conf_scores,
- function_select='integral')
-
- all_ap.append(ap)
-
- # depth_map, _ = visualize_depth_numpy(depth_map.numpy(),near_far)
- # if len(test_dataset.all_rgbs):
- # gt_rgb = test_dataset.all_rgbs[idxs[idx]].view(H, W, 3)
- # loss = torch.mean((rgb_map - gt_rgb) ** 2)
- # PSNRs.append(-10.0 * np.log(loss.item()) / np.log(10.0))
- #
- # if compute_extra_metrics:
- # ssim = rgb_ssim(rgb_map, gt_rgb, 1)
- # l_a = rgb_lpips(gt_rgb.numpy(), rgb_map.numpy(), 'alex', sem_tensorf.device)
- # l_v = rgb_lpips(gt_rgb.numpy(), rgb_map.numpy(), 'vgg', sem_tensorf.device)
- # ssims.append(ssim)
- # l_alex.append(l_a)
- # l_vgg.append(l_v)
-
- # rgb_map = (rgb_map.numpy() * 255).astype('uint8')
- # rgb_map = np.concatenate((rgb_map, depth_map), axis=1)
- # rgb_maps.append(rgb_map)
- # sem_logit_map = sem_logit_map.reshape(H, W, num_semantic_class)
- sem_label_map = sem_label_map.reshape(H, W).astype('uint8')
- sem_map = sem_map.reshape(H, W, 3).astype('uint8')
-
- sem_maps.append(sem_map)
- sem_label_maps.append(sem_label_map)
- # depth_maps.append(depth_map)
- if savePath is not None:
- # imageio.imwrite(f'{savePath}/{prtx}{idx:03d}.png', rgb_map)
- # rgb_map = np.concatenate((rgb_map, depth_map), axis=1)
- # imageio.imwrite(f'{savePath}/rgbd/{prtx}{idx:03d}.png', rgb_map)
- imageio.imwrite(f'{savePath}/sem/{prtx}{idx:03d}.png', sem_map)
- imageio.imwrite(f'{savePath}/sem/{prtx}label_{idx:03d}.png', sem_label_map)
- # np.save(f'{savePath}/sem/{prtx}logit_{idx:03d}.npy', sem_logit_map)
- # np.savez_compressed(f'{savePath}/sem/{prtx}logit_{idx:03d}.npz', sem_logit_map)
-
- miou = np.mean(np.array(metric_ious))
-
- all_ap = np.array(all_ap)
- mean_ap = np.mean(all_ap, axis=0)
-
- sem_label_maps = np.stack(sem_label_maps, 0)
- sem_label_maps_flatten = sem_label_maps.reshape(-1, 1)
- gt_label_maps_flatten = gt_label_maps[0::img_eval_interval].reshape(-1, 1)
- metric_iou_pix = val_cm_pix.add_batch(sem_label_maps_flatten, gt_label_maps_flatten, return_miou=True)
- # metric_ious.append(metric_iou_pix)
-
- sparse_ratio = 1 - 1 / args.sem_interval
- ins_metric_desc = f'sparsity_ratio = {sparse_ratio * 100}%, mIoU = {miou}, '
- for i, thre in enumerate(thre_list):
- ins_metric_desc += f'mAP@{thre} = {mean_ap[i]}, '
- ins_metric_desc += f'mIoU_pix = {metric_iou_pix}'
-
- # imageio.mimwrite(f'{savePath}/{prtx}video.mp4', np.stack(rgb_maps), fps=30, quality=10)
- # imageio.mimwrite(f'{savePath}/{prtx}depthvideo.mp4', np.stack(depth_maps), fps=30, quality=10)
- imageio.mimwrite(f'{savePath}/{prtx}sem_video.mp4', np.stack(sem_maps), fps=10, quality=10)
- with open(f'{savePath}/{prtx}ins_metrics.txt', 'w') as f:
- f.write(ins_metric_desc)
-
- # if PSNRs:
- # psnr = np.mean(np.asarray(PSNRs))
- # if compute_extra_metrics:
- # ssim = np.mean(np.asarray(ssims))
- # l_a = np.mean(np.asarray(l_alex))
- # l_v = np.mean(np.asarray(l_vgg))
- # np.savetxt(f'{savePath}/{prtx}mean.txt', np.asarray([psnr, ssim, l_a, l_v]))
- # else:
- # np.savetxt(f'{savePath}/{prtx}mean.txt', np.asarray([psnr]))
- #
- # return PSNRs
-
-
- @torch.no_grad()
- def evaluation_path(test_dataset, tensorf, c2ws, renderer, savePath=None, N_vis=5, prtx='', N_samples=-1,
- white_bg=False, ndc_ray=False, compute_extra_metrics=True, device='cuda'):
- PSNRs, rgb_maps, depth_maps = [], [], []
- ssims,l_alex,l_vgg=[],[],[]
- os.makedirs(savePath, exist_ok=True)
- os.makedirs(savePath+"/rgbd", exist_ok=True)
-
- try:
- tqdm._instances.clear()
- except Exception:
- pass
-
- near_far = test_dataset.near_far
- for idx, c2w in tqdm(enumerate(c2ws)):
- W, H = test_dataset.img_wh
-
- c2w = torch.FloatTensor(c2w)
- rays_o, rays_d = get_rays(test_dataset.directions, c2w) # both (h*w, 3)
- if ndc_ray:
- rays_o, rays_d = ndc_rays(H, W, test_dataset.focal[0], 1.0, rays_o, rays_d)
- rays = torch.cat([rays_o, rays_d], 1) # (h*w, 6)
-
- rgb_map, _, depth_map, _, _ = renderer(rays, tensorf, chunk=8192, N_samples=N_samples,
- ndc_ray=ndc_ray, white_bg = white_bg, device=device)
- rgb_map = rgb_map.clamp(0.0, 1.0)
-
- rgb_map, depth_map = rgb_map.reshape(H, W, 3).cpu(), depth_map.reshape(H, W).cpu()
-
- depth_map, _ = visualize_depth_numpy(depth_map.numpy(),near_far)
-
- rgb_map = (rgb_map.numpy() * 255).astype('uint8')
- # rgb_map = np.concatenate((rgb_map, depth_map), axis=1)
- rgb_maps.append(rgb_map)
- depth_maps.append(depth_map)
- if savePath is not None:
- imageio.imwrite(f'{savePath}/{prtx}{idx:03d}.png', rgb_map)
- rgb_map = np.concatenate((rgb_map, depth_map), axis=1)
- imageio.imwrite(f'{savePath}/rgbd/{prtx}{idx:03d}.png', rgb_map)
-
- imageio.mimwrite(f'{savePath}/{prtx}video.mp4', np.stack(rgb_maps), fps=30, quality=8)
- imageio.mimwrite(f'{savePath}/{prtx}depthvideo.mp4', np.stack(depth_maps), fps=30, quality=8)
-
- if PSNRs:
- psnr = np.mean(np.asarray(PSNRs))
- if compute_extra_metrics:
- ssim = np.mean(np.asarray(ssims))
- l_a = np.mean(np.asarray(l_alex))
- l_v = np.mean(np.asarray(l_vgg))
- np.savetxt(f'{savePath}/{prtx}mean.txt', np.asarray([psnr, ssim, l_a, l_v]))
- else:
- np.savetxt(f'{savePath}/{prtx}mean.txt', np.asarray([psnr]))
-
- return PSNRs
-
-
- @torch.no_grad()
- def evaluation_path_with_sem(test_dataset, sem_tensorf, c2ws, sem_renderer, savePath=None, N_vis=5, prtx='', N_samples=-1,
- white_bg=False, ndc_ray=False, compute_extra_metrics=True, fp16=False, device='cuda',
- chunk_size=2048):
- """
- Render images from new view-points
- """
- PSNRs, rgb_maps, sem_maps, depth_maps = [], [], [], []
- ssims, l_alex, l_vgg = [], [], []
- os.makedirs(savePath, exist_ok=True)
- os.makedirs(savePath + "/rgbd", exist_ok=True)
- os.makedirs(savePath + "/sem", exist_ok=True)
-
- try:
- tqdm._instances.clear()
- except Exception:
- pass
-
- near_far = test_dataset.near_far
- for idx, c2w in tqdm(enumerate(c2ws)):
- W, H = test_dataset.img_wh
-
- c2w = torch.FloatTensor(c2w)
- rays_o, rays_d = get_rays(test_dataset.directions, c2w) # both (h*w, 3)
- if ndc_ray:
- rays_o, rays_d = ndc_rays(H, W, test_dataset.focal[0], 1.0, rays_o, rays_d)
- rays = torch.cat([rays_o, rays_d], 1) # (h*w, 6)
-
- rgb_map, _, sem_map, semfea_map, depth_map, _, _, loss_3d = sem_renderer(rays, sem_tensorf, chunk=chunk_size,
- N_samples=N_samples, # changeable chunk_size
- ndc_ray=ndc_ray, white_bg=white_bg, fp16=fp16,
- device=device)
- rgb_map = rgb_map.clamp(0.0, 1.0)
-
- rgb_map, depth_map = rgb_map.reshape(H, W, 3).cpu(), depth_map.reshape(H, W).cpu()
-
- depth_map, _ = visualize_depth_numpy(depth_map.numpy(),near_far)
-
- rgb_map = (rgb_map.numpy() * 255).astype('uint8')
- # rgb_map = np.concatenate((rgb_map, depth_map), axis=1)
- rgb_maps.append(rgb_map)
-
- # sem_map = sem_map.clamp(0.0, 1.0)
- logits_2_label = lambda x: torch.argmax(torch.nn.functional.softmax(x, dim=-1), dim=-1)
- sem_map = logits_2_label(sem_map)
- sem_map = sem_map.cpu().numpy()
- sem_map_label = sem_map.copy() # ?
- sem_map = test_dataset.label_color_map(sem_map)
- # ignore: black
- ignore_label = -1
- sem_map[sem_map_label == ignore_label, :] = 0
-
- sem_map = sem_map.reshape(H, W, 3).astype('uint8')
- sem_maps.append(sem_map)
-
- depth_maps.append(depth_map)
- if savePath is not None:
- imageio.imwrite(f'{savePath}/{prtx}{idx:03d}.png', rgb_map)
- rgb_map = np.concatenate((rgb_map, depth_map), axis=1)
- imageio.imwrite(f'{savePath}/rgbd/{prtx}{idx:03d}.png', rgb_map)
- imageio.imwrite(f'{savePath}/sem/{prtx}{idx:03d}.png', sem_map)
-
- imageio.mimwrite(f'{savePath}/{prtx}video.mp4', np.stack(rgb_maps), fps=30, quality=8)
- imageio.mimwrite(f'{savePath}/{prtx}depthvideo.mp4', np.stack(depth_maps), fps=30, quality=8)
- imageio.mimwrite(f'{savePath}/{prtx}semvideo.mp4', np.stack(sem_maps), fps=30, quality=8)
-
- if PSNRs:
- psnr = np.mean(np.asarray(PSNRs))
- if compute_extra_metrics:
- ssim = np.mean(np.asarray(ssims))
- l_a = np.mean(np.asarray(l_alex))
- l_v = np.mean(np.asarray(l_vgg))
- np.savetxt(f'{savePath}/{prtx}mean.txt', np.asarray([psnr, ssim, l_a, l_v]))
- else:
- np.savetxt(f'{savePath}/{prtx}mean.txt', np.asarray([psnr]))
-
- return PSNRs
-
-
- @torch.no_grad()
- def tps_rendering(dataset, sem_tensorf, args, sem_renderer, savePath=None, N_vis=5, prtx='', N_samples=-1,
- white_bg=False, ndc_ray=False, compute_extra_metrics=True, device='cuda',
- tps_function=None,
- chunk_size=4096):
- rgb_maps, depth_maps, sem_maps = [], [], []
- # ssims,l_alex,l_vgg=[],[],[]
- os.makedirs(savePath, exist_ok=True)
- os.makedirs(savePath+"/rgb", exist_ok=True)
- os.makedirs(savePath+"/sem", exist_ok=True)
- os.makedirs(savePath+"/semantic_instance", exist_ok=True)
-
- os.makedirs(savePath+"/rgbd", exist_ok=True)
-
- try:
- tqdm._instances.clear()
- except Exception:
- pass
-
- near_far = dataset.near_far
- img_eval_interval = 1 if N_vis < 0 else max(dataset.all_rays.shape[0] // N_vis, 1)
- idxs = list(range(0, dataset.all_rays.shape[0], img_eval_interval))
-
- for idx, samples in tqdm(enumerate(dataset.all_rays[0::img_eval_interval]), file=sys.stdout):
- W, H = dataset.img_wh
- rays = samples.view(-1, samples.shape[-1]) # rays for one image
-
- rgb_map, _, sem_map, depth_map, _, _,loss_3d = sem_renderer(rays, sem_tensorf, chunk=chunk_size, N_samples=N_samples, # sem_map: (chunk_size, num_sem_classes)
- ndc_ray=ndc_ray, white_bg=white_bg, device=device,
- tps_function=tps_function)
-
- rgb_map = rgb_map.clamp(0.0, 1.0) # (H*W, 3)
- rgb_map, depth_map = rgb_map.reshape(H, W, 3).cpu(), depth_map.reshape(H, W).cpu()
-
- logits_2_label = lambda x: torch.argmax(torch.nn.functional.softmax(x, dim=-1), dim=-1)
- sem_map = logits_2_label(sem_map) # (chunk_size, )
- sem_map = sem_map.cpu().numpy()
-
- sem_map_pseudo_label = dataset.inv_map_sem_gt_label(sem_map)
- sem_map_pseudo_label = sem_map_pseudo_label.reshape(H, W).astype('uint8')
- imageio.imwrite(f'{savePath}/semantic_instance/{prtx}semantic_instance_{idx}.png', sem_map_pseudo_label)
-
- sem_map_label = sem_map.copy() # ?
- """not the same as replica (() or [])"""
- sem_map = dataset.label_color_map(sem_map) # ndarray ?
- if args.dataset_name == "replica":
- # used in replica ?
- ignore_label = -1
- sem_map[sem_map_label == ignore_label, :] = 0
-
- sem_map = sem_map.reshape(H, W, 3).astype('uint8')
-
- depth_map, _ = visualize_depth_numpy(depth_map.numpy(), near_far)
-
- rgb_map = (rgb_map.numpy() * 255).astype('uint8')
- # rgb_map = np.concatenate((rgb_map, depth_map), axis=1)
- rgb_maps.append(rgb_map)
- sem_maps.append(sem_map)
- depth_maps.append(depth_map)
- if savePath is not None:
- imageio.imwrite(f'{savePath}/rgb/{prtx}rgb_{idx}.png', rgb_map)
- imageio.imwrite(f'{savePath}/sem/{prtx}{idx:03d}.png', sem_map)
-
- rgb_map = np.concatenate((rgb_map, depth_map), axis=1)
- imageio.imwrite(f'{savePath}/rgbd/{prtx}{idx:03d}.png', rgb_map)
-
- imageio.mimwrite(f'{savePath}/{prtx}video.mp4', np.stack(rgb_maps), fps=30, quality=10)
- imageio.mimwrite(f'{savePath}/{prtx}depthvideo.mp4', np.stack(depth_maps), fps=30, quality=10)
- imageio.mimwrite(f'{savePath}/{prtx}sem_video.mp4', np.stack(sem_maps), fps=10, quality=10)
|