renderer.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys,os,imageio,lpips\n",
    "root = '/home/hengfei/Desktop/research/mvsnerf'\n",
    "os.chdir(root)\n",
    "sys.path.append(root)\n",
    "\n",
    "from opt import config_parser\n",
    "from data import dataset_dict\n",
    "from torch.utils.data import DataLoader\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "\n",
    "# models\n",
    "from models import *\n",
    "from renderer import *\n",
    "from data.ray_utils import get_rays\n",
    "\n",
    "from tqdm import tqdm\n",
    "\n",
    "\n",
    "from skimage.metrics import structural_similarity\n",
    "\n",
    "# pytorch-lightning\n",
    "from pytorch_lightning.callbacks import ModelCheckpoint\n",
    "from pytorch_lightning import LightningModule, Trainer, loggers\n",
    "\n",
    "\n",
    "from data.ray_utils import ray_marcher\n",
    "\n",
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "\n",
    "torch.cuda.set_device(1)\n",
    "os.environ['CUDA_VISIBLE_DEVICES'] = '1'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Setting up [LPIPS] perceptual loss: trunk [vgg], v[0.1], spatial [off]\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Downloading: \"https://download.pytorch.org/models/vgg16-397923af.pth\" to /home/hengfei/.cache/torch/hub/checkpoints/vgg16-397923af.pth\n",
      "100%|██████████| 528M/528M [02:05<00:00, 4.42MB/s] \n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loading model from: /home/hengfei/miniconda3/envs/mvsnerf/lib/python3.8/site-packages/lpips/weights/v0.1/vgg.pth\n"
     ]
    }
   ],
   "source": [
    "def decode_batch(batch):\n",
    "    rays = batch['rays']  # (B, 8)\n",
    "    rgbs = batch['rgbs']  # (B, 3)\n",
    "    return rays, rgbs\n",
    "\n",
    "def unpreprocess(data, shape=(1,1,3,1,1)):\n",
    "    # to unnormalize image for visualization\n",
    "    # data N V C H W\n",
    "    device = data.device\n",
    "    mean = torch.tensor([-0.485 / 0.229, -0.456 / 0.224, -0.406 / 0.225]).view(*shape).to(device)\n",
    "    std = torch.tensor([1 / 0.229, 1 / 0.224, 1 / 0.225]).view(*shape).to(device)\n",
    "\n",
    "    return (data - mean) / std\n",
    "\n",
    "def read_depth(filename):\n",
    "    depth_h = np.array(read_pfm(filename)[0], dtype=np.float32) # (800, 800)\n",
    "    depth_h = cv2.resize(depth_h, None, fx=0.5, fy=0.5,\n",
    "                       interpolation=cv2.INTER_NEAREST)  # (600, 800)\n",
    "    depth_h = depth_h[44:556, 80:720]  # (512, 640)\n",
    "#     depth = cv2.resize(depth_h, None, fx=0.5, fy=0.5,interpolation=cv2.INTER_NEAREST)#!!!!!!!!!!!!!!!!!!!!!!!!!\n",
    "    mask = depth>0\n",
    "    return depth_h,mask\n",
    "\n",
    "loss_fn_vgg = lpips.LPIPS(net='vgg') \n",
    "mse2psnr = lambda x : -10. * np.log(x) / np.log(10.)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# llff no fine tuning"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## rendering novel views with nearest 3 source views"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Found ckpts ['./ckpts/mvsnerf-v0.tar']\n",
      "Reloading from ./ckpts/mvsnerf-v0.tar\n",
      "============> rendering dataset <===================\n",
      "11 11 /home/hengfei/Desktop/research/mvsnerf/xgaze/xgaze_11images_cropped_colmapCODE\n",
      "===> valing index: [0 5]\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/hengfei/miniconda3/envs/mvsnerf/lib/python3.8/site-packages/torch/functional.py:445: UserWarning: torch.meshgrid: in an upcoming release, it will be required to pass the indexing argument. (Triggered internally at  ../aten/src/ATen/native/TensorShape.cpp:2157.)\n",
      "  return _VF.meshgrid(tensors, **kwargs)  # type: ignore[attr-defined]\n",
      "  0%|          | 0/2 [00:00<?, ?it/s]/tmp/ipykernel_7078/3929512484.py:111: FutureWarning: `multichannel` is a deprecated argument name for `structural_similarity`. It will be removed in version 1.0.Please use `channel_axis` instead.\n",
      "  ssim.append( structural_similarity(rgb_rays, img, multichannel=True))\n",
      "100%|██████████| 2/2 [00:25<00:00, 12.72s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "=====> scene: xgaze_11images_cropped_colmapCODE mean psnr 20.593928452628838 ssim: 0.8184963464736938 lpips: 0.5154071301221848\n",
      "=====> all mean psnr 20.593928452628838 ssim: 0.8184963464736938 lpips: 0.5154071301221848\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "psnr_all,ssim_all,LPIPS_vgg_all = [],[],[]\n",
    "for i_scene, scene in enumerate([ 'xgaze_11images_cropped_colmapCODE']):#\n",
    "    psnr,ssim,LPIPS_vgg = [],[],[]\n",
    "    cmd = f'--datadir /home/hengfei/Desktop/research/mvsnerf/xgaze/{scene}  \\\n",
    "     --dataset_name llff \\\n",
    "     --net_type v0 --ckpt ./ckpts/mvsnerf-v0.tar '\n",
    "\n",
    "\n",
    "    args = config_parser(cmd.split())\n",
    "    args.use_viewdirs = True\n",
    "\n",
    "    args.N_samples = 128\n",
    "    args.feat_dim = 8+12\n",
    "\n",
    "    # create models\n",
    "    if 0==i_scene:\n",
    "        render_kwargs_train, render_kwargs_test, start, grad_vars = create_nerf_mvs(args, use_mvs=True, dir_embedder=False, pts_embedder=True)\n",
    "        filter_keys(render_kwargs_train)\n",
    "\n",
    "        MVSNet = render_kwargs_train['network_mvs']\n",
    "        render_kwargs_train.pop('network_mvs')\n",
    "\n",
    "\n",
    "    datadir = args.datadir\n",
    "    datatype = 'val'\n",
    "    pad = 24\n",
    "    args.chunk = 5120\n",
    "\n",
    "\n",
    "    print('============> rendering dataset <===================')\n",
    "    dataset = dataset_dict[args.dataset_name](args, split=datatype)\n",
    "    val_idx = dataset.img_idx\n",
    "    \n",
    "    save_as_image = True\n",
    "    save_dir = f'results/xgaze_11images_cropped_colmapCODE'\n",
    "    os.makedirs(save_dir, exist_ok=True)\n",
    "    MVSNet.train()\n",
    "    MVSNet = MVSNet.cuda()\n",
    "    \n",
    "    with torch.no_grad():\n",
    "        \n",
    "        try:\n",
    "            tqdm._instances.clear() \n",
    "        except Exception:     \n",
    "            pass\n",
    "        \n",
    "        for i, batch in enumerate(tqdm(dataset)):\n",
    "            torch.cuda.empty_cache()\n",
    "\n",
    "                \n",
    "            rays, img = decode_batch(batch)\n",
    "            rays = rays.squeeze().to(device)  # (H*W, 3)\n",
    "            img = img.squeeze().cpu().numpy()  # (H, W, 3)\n",
    "        \n",
    "            # find nearest image idx\n",
    "            img_idx = torch.load('configs/pairs.th')[f'{scene}_train']\n",
    "            positions = dataset.poses[img_idx,:3,3]\n",
    "            dis = np.sum(np.abs(positions - dataset.poses[val_idx[i],:3,3]), axis=-1)\n",
    "            pair_idx = np.argsort(dis)[:3]\n",
    "            pair_idx = [img_idx[item] for item in pair_idx]\n",
    "            \n",
    "            imgs_source, proj_mats, near_far_source, pose_source = dataset.read_source_views(pair_idx=pair_idx,device=device)\n",
    "            volume_feature, img_feat, _ = MVSNet(imgs_source, proj_mats, near_far_source, pad=pad, lindisp=False)\n",
    "            imgs_source = unpreprocess(imgs_source)\n",
    "        \n",
    "            N_rays_all = rays.shape[0]\n",
    "            rgb_rays, depth_rays_preds = [],[]\n",
    "            for chunk_idx in range(N_rays_all//args.chunk + int(N_rays_all%args.chunk>0)):\n",
    "\n",
    "                xyz_coarse_sampled, rays_o, rays_d, z_vals = ray_marcher(rays[chunk_idx*args.chunk:(chunk_idx+1)*args.chunk],\n",
    "                                                    N_samples=args.N_samples, lindisp=False)\n",
    "\n",
    "                # Converting world coordinate to ndc coordinate\n",
    "                H, W = img.shape[:2]\n",
    "                inv_scale = torch.tensor([W - 1, H - 1]).to(device)\n",
    "                w2c_ref, intrinsic_ref = pose_source['w2cs'][0], pose_source['intrinsics'][0].clone()\n",
    "                xyz_NDC = get_ndc_coordinate(w2c_ref, intrinsic_ref, xyz_coarse_sampled, inv_scale,\n",
    "                        near=near_far_source[0], far=near_far_source[1], pad=pad*args.imgScale_test, lindisp=False)\n",
    "\n",
    "\n",
    "                # rendering\n",
    "                rgb, disp, acc, depth_pred, alpha, extras = rendering(args, pose_source, xyz_coarse_sampled,\n",
    "                    xyz_NDC, z_vals, rays_o, rays_d,volume_feature,imgs_source, img_feat=None,**render_kwargs_train)\n",
    "    \n",
    "                \n",
    "                rgb, depth_pred = torch.clamp(rgb.cpu(),0,1.0).numpy(), depth_pred.cpu().numpy()\n",
    "                rgb_rays.append(rgb)\n",
    "                depth_rays_preds.append(depth_pred)\n",
    "\n",
    "\n",
    "            depth_rays_preds = np.concatenate(depth_rays_preds).reshape(H, W)\n",
    "            depth_rays_preds, _ = visualize_depth_numpy(depth_rays_preds, near_far_source)\n",
    "            \n",
    "            rgb_rays = np.concatenate(rgb_rays).reshape(H, W, 3)\n",
    "            img_vis = np.concatenate((img*255,rgb_rays*255,depth_rays_preds),axis=1)\n",
    "            \n",
    "#             img_vis = np.concatenate((torch.cat(torch.split(imgs_source*255, [1,1,1], dim=1),-1).squeeze().permute(1,2,0).cpu().numpy(),img_vis),axis=1)\n",
    "            \n",
    "            if save_as_image:\n",
    "                imageio.imwrite(f'{save_dir}/{scene}_{val_idx[i]:03d}.png', img_vis.astype('uint8'))\n",
    "            else:\n",
    "                rgbs.append(img_vis.astype('uint8'))\n",
    "                \n",
    "            # quantity\n",
    "            # center crop 0.8 ratio\n",
    "            H_crop, W_crop = np.array(rgb_rays.shape[:2])//10\n",
    "            img = img[H_crop:-H_crop,W_crop:-W_crop]\n",
    "            rgb_rays = rgb_rays[H_crop:-H_crop,W_crop:-W_crop]\n",
    "            \n",
    "            psnr.append( mse2psnr(np.mean((rgb_rays-img)**2)))\n",
    "            ssim.append( structural_similarity(rgb_rays, img, multichannel=True))\n",
    "            \n",
    "            img_tensor = torch.from_numpy(rgb_rays)[None].permute(0,3,1,2).float()*2-1.0 # image should be RGB, IMPORTANT: normalized to [-1,1]\n",
    "            img_gt_tensor = torch.from_numpy(img)[None].permute(0,3,1,2).float()*2-1.0\n",
    "            LPIPS_vgg.append( loss_fn_vgg(img_tensor, img_gt_tensor).item())\n",
    "            \n",
    "        print(f'=====> scene: {scene} mean psnr {np.mean(psnr)} ssim: {np.mean(ssim)} lpips: {np.mean(LPIPS_vgg)}')   \n",
    "        psnr_all.append(psnr);ssim_all.append(ssim);LPIPS_vgg_all.append(LPIPS_vgg)\n",
    "    \n",
    "    if not save_as_image:\n",
    "        imageio.mimwrite(f'{save_dir}/{scene}_spiral.mp4', np.stack(rgbs), fps=20, quality=10)\n",
    "print(f'=====> all mean psnr {np.mean(psnr_all)} ssim: {np.mean(ssim_all)} lpips: {np.mean(LPIPS_vgg_all)}') "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### for room"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "psnr_all,ssim_all,LPIPS_vgg_all = [],[],[]\n",
    "for i_scene, scene in enumerate([ 'room']):#\n",
    "    psnr,ssim,LPIPS_vgg = [],[],[]\n",
    "    cmd = f'--datadir /mnt/new_disk_2/anpei/Dataset/MVSNeRF/nerf_llff_data/{scene}  \\\n",
    "     --dataset_name llff \\\n",
    "     --net_type v0 --ckpt ./ckpts/mvsnerf-v0.tar '\n",
    "\n",
    "\n",
    "    args = config_parser(cmd.split())\n",
    "    args.use_viewdirs = True\n",
    "\n",
    "    args.N_samples = 128\n",
    "    args.feat_dim = 8+12\n",
    "\n",
    "    # create models\n",
    "    if 0==i_scene:\n",
    "        render_kwargs_train, render_kwargs_test, start, grad_vars = create_nerf_mvs(args, use_mvs=True, dir_embedder=False, pts_embedder=True)\n",
    "        filter_keys(render_kwargs_train)\n",
    "\n",
    "        MVSNet = render_kwargs_train['network_mvs']\n",
    "        render_kwargs_train.pop('network_mvs')\n",
    "\n",
    "\n",
    "    datadir = args.datadir\n",
    "    datatype = 'val'\n",
    "    pad = 24\n",
    "    args.chunk = 5120\n",
    "\n",
    "\n",
    "    print('============> rendering dataset <===================')\n",
    "    dataset = dataset_dict[args.dataset_name](args, split=datatype)\n",
    "    val_idx = dataset.img_idx\n",
    "    \n",
    "    save_as_image = True\n",
    "    save_dir = f'results/test3'\n",
    "    os.makedirs(save_dir, exist_ok=True)\n",
    "    MVSNet.train()\n",
    "    MVSNet = MVSNet.cuda()\n",
    "    \n",
    "    with torch.no_grad():\n",
    "        \n",
    "        try:\n",
    "            tqdm._instances.clear() \n",
    "        except Exception:     \n",
    "            pass\n",
    "        \n",
    "        for i, batch in enumerate(tqdm(dataset)):\n",
    "            torch.cuda.empty_cache()\n",
    "            if 1!=i:\n",
    "                continue\n",
    "                \n",
    "            rays, img = decode_batch(batch)\n",
    "            rays = rays.squeeze().to(device)  # (H*W, 3)\n",
    "            img = img.squeeze().cpu().numpy()  # (H, W, 3)\n",
    "        \n",
    "            # find nearest image idx\n",
    "            img_idx = torch.load('configs/pairs.th')[f'{scene}_train']\n",
    "            positions = dataset.poses[img_idx,:3,3]\n",
    "            dis = np.sum(np.abs(positions - dataset.poses[val_idx[i],:3,3]), axis=-1)\n",
    "            pair_idx = np.argsort(dis)[:3]\n",
    "            pair_idx = img_idx[:3]#[img_idx[item] for item in pair_idx]\n",
    "            \n",
    "            imgs_source, proj_mats, near_far_source, pose_source = dataset.read_source_views(pair_idx=pair_idx,device=device)\n",
    "            volume_feature, img_feat, _ = MVSNet(imgs_source, proj_mats, near_far_source, pad=pad, lindisp=False)\n",
    "            imgs_source = unpreprocess(imgs_source)\n",
    "        \n",
    "            N_rays_all = rays.shape[0]\n",
    "            rgb_rays, depth_rays_preds = [],[]\n",
    "            for chunk_idx in range(N_rays_all//args.chunk + int(N_rays_all%args.chunk>0)):\n",
    "\n",
    "                xyz_coarse_sampled, rays_o, rays_d, z_vals = ray_marcher(rays[chunk_idx*args.chunk:(chunk_idx+1)*args.chunk],\n",
    "                                                    N_samples=args.N_samples, lindisp=False)\n",
    "\n",
    "                # Converting world coordinate to ndc coordinate\n",
    "                H, W = img.shape[:2]\n",
    "                inv_scale = torch.tensor([W - 1, H - 1]).to(device)\n",
    "                w2c_ref, intrinsic_ref = pose_source['w2cs'][0], pose_source['intrinsics'][0].clone()\n",
    "                xyz_NDC = get_ndc_coordinate(w2c_ref, intrinsic_ref, xyz_coarse_sampled, inv_scale,\n",
    "                        near=near_far_source[0], far=near_far_source[1], pad=pad*args.imgScale_test, lindisp=False)\n",
    "\n",
    "\n",
    "                # rendering\n",
    "                rgb, disp, acc, depth_pred, alpha, extras = rendering(args, pose_source, xyz_coarse_sampled,\n",
    "                    xyz_NDC, z_vals, rays_o, rays_d,volume_feature,imgs_source, img_feat=None,**render_kwargs_train)\n",
    "    \n",
    "                \n",
    "                rgb, depth_pred = torch.clamp(rgb.cpu(),0,1.0).numpy(), depth_pred.cpu().numpy()\n",
    "                rgb_rays.append(rgb)\n",
    "                depth_rays_preds.append(depth_pred)\n",
    "\n",
    "\n",
    "            depth_rays_preds = np.concatenate(depth_rays_preds).reshape(H, W)\n",
    "            depth_rays_preds, _ = visualize_depth_numpy(depth_rays_preds, near_far_source)\n",
    "            \n",
    "            rgb_rays = np.concatenate(rgb_rays).reshape(H, W, 3)\n",
    "            img_vis = np.concatenate((img*255,rgb_rays*255,depth_rays_preds),axis=1)\n",
    "            \n",
    "#             img_vis = np.concatenate((torch.cat(torch.split(imgs_source*255, [1,1,1], dim=1),-1).squeeze().permute(1,2,0).cpu().numpy(),img_vis),axis=1)\n",
    "            \n",
    "            if save_as_image:\n",
    "                imageio.imwrite(f'{save_dir}/{scene}_{val_idx[i]:03d}.png', img_vis.astype('uint8'))\n",
    "            else:\n",
    "                rgbs.append(img_vis.astype('uint8'))\n",
    "                \n",
    "            # quantity\n",
    "            # center crop 0.8 ratio\n",
    "            H_crop, W_crop = np.array(rgb_rays.shape[:2])//10\n",
    "            img = img[H_crop:-H_crop,W_crop:-W_crop]\n",
    "            rgb_rays = rgb_rays[H_crop:-H_crop,W_crop:-W_crop]\n",
    "            \n",
    "            psnr.append( mse2psnr(np.mean((rgb_rays-img)**2)))\n",
    "            ssim.append( structural_similarity(rgb_rays, img, multichannel=True))\n",
    "            \n",
    "            img_tensor = torch.from_numpy(rgb_rays)[None].permute(0,3,1,2).float()*2-1.0 # image should be RGB, IMPORTANT: normalized to [-1,1]\n",
    "            img_gt_tensor = torch.from_numpy(img)[None].permute(0,3,1,2).float()*2-1.0\n",
    "            LPIPS_vgg.append( loss_fn_vgg(img_tensor, img_gt_tensor).item())\n",
    "            \n",
    "        print(f'=====> scene: {scene} mean psnr {np.mean(psnr)} ssim: {np.mean(ssim)} lpips: {np.mean(LPIPS_vgg)}')   \n",
    "        psnr_all.append(psnr);ssim_all.append(ssim);LPIPS_vgg_all.append(LPIPS_vgg)\n",
    "    \n",
    "    if not save_as_image:\n",
    "        imageio.mimwrite(f'{save_dir}/{scene}_spiral.mp4', np.stack(rgbs), fps=20, quality=10)\n",
    "print(f'=====> all mean psnr {np.mean(psnr_all)} ssim: {np.mean(ssim_all)} lpips: {np.mean(LPIPS_vgg_all)}') "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## rendering novel views with fixed 3 source views"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "psnr_all,ssim_all,LPIPS_vgg_all = [],[],[]\n",
    "for i_scene, scene in enumerate(['fortress','flower','orchids', 'room','leaves','horns','trex','fern']):#'flower','orchids', 'room','leaves','fern','horns','trex','fortress'\n",
    "    psnr,ssim,LPIPS_vgg = [],[],[]\n",
    "    cmd = f'--datadir /mnt/new_disk_2/anpei/Dataset/MVSNeRF/nerf_llff_data/{scene}  \\\n",
    "     --dataset_name llff \\\n",
    "     --ckpt ./ckpts/mvsnerf-v0.tar  \\\n",
    "     --net_type v0 --netwidth 128 --netdepth 6'\n",
    "\n",
    "    args = config_parser(cmd.split())\n",
    "    args.use_viewdirs = True\n",
    "\n",
    "    args.N_samples = 128\n",
    "    args.feat_dim =  8+3*4\n",
    "\n",
    "    # create models\n",
    "    if 0==i_scene:\n",
    "        render_kwargs_train, render_kwargs_test, start, grad_vars = create_nerf_mvs(args, use_mvs=True, dir_embedder=False, pts_embedder=True)\n",
    "        filter_keys(render_kwargs_train)\n",
    "\n",
    "        MVSNet = render_kwargs_train['network_mvs']\n",
    "        render_kwargs_train.pop('network_mvs')\n",
    "\n",
    "\n",
    "    datadir = args.datadir\n",
    "    datatype = 'val'\n",
    "    pad = 24\n",
    "    args.chunk = 5120\n",
    "\n",
    "\n",
    "    print('============> rendering dataset <===================')\n",
    "    dataset = dataset_dict[args.dataset_name](args, split=datatype)\n",
    "    val_idx = dataset.img_idx\n",
    "    \n",
    "    save_as_image = True\n",
    "    save_dir = f'results/test3'\n",
    "    os.makedirs(save_dir, exist_ok=True)\n",
    "    MVSNet.train()\n",
    "    MVSNet = MVSNet.cuda()\n",
    "    \n",
    "    with torch.no_grad():\n",
    "\n",
    "        imgs_source, proj_mats, near_far_source, pose_source = dataset.read_source_views(device=device)\n",
    "        volume_feature, _, _ = MVSNet(imgs_source, proj_mats, near_far_source, pad=pad)\n",
    "        imgs_source = unpreprocess(imgs_source)\n",
    "\n",
    "        try:\n",
    "            tqdm._instances.clear() \n",
    "        except Exception:     \n",
    "            pass\n",
    "        \n",
    "        for i, batch in enumerate(tqdm(dataset)):\n",
    "            torch.cuda.empty_cache()\n",
    "            \n",
    "            rays, img = decode_batch(batch)\n",
    "            rays = rays.squeeze().to(device)  # (H*W, 3)\n",
    "            img = img.squeeze().cpu().numpy()  # (H, W, 3)\n",
    "        \n",
    "            N_rays_all = rays.shape[0]\n",
    "            rgb_rays, depth_rays_preds = [],[]\n",
    "            for chunk_idx in range(N_rays_all//args.chunk + int(N_rays_all%args.chunk>0)):\n",
    "\n",
    "                xyz_coarse_sampled, rays_o, rays_d, z_vals = ray_marcher(rays[chunk_idx*args.chunk:(chunk_idx+1)*args.chunk],\n",
    "                                                    N_samples=args.N_samples)\n",
    "\n",
    "                # Converting world coordinate to ndc coordinate\n",
    "                H, W = img.shape[:2]\n",
    "                inv_scale = torch.tensor([W - 1, H - 1]).to(device)\n",
    "                w2c_ref, intrinsic_ref = pose_source['w2cs'][0], pose_source['intrinsics'][0].clone()\n",
    "                xyz_NDC = get_ndc_coordinate(w2c_ref, intrinsic_ref, xyz_coarse_sampled, inv_scale,\n",
    "                                             near=near_far_source[0], far=near_far_source[1], pad=pad*args.imgScale_test)\n",
    "\n",
    "\n",
    "                # rendering\n",
    "                rgb, disp, acc, depth_pred, alpha, extras = rendering(args, pose_source, xyz_coarse_sampled,\n",
    "                                                                       xyz_NDC, z_vals, rays_o, rays_d,\n",
    "                                                                       volume_feature,imgs_source, **render_kwargs_train)\n",
    "    \n",
    "                \n",
    "                rgb, depth_pred = torch.clamp(rgb.cpu(),0,1.0).numpy(), depth_pred.cpu().numpy()\n",
    "                rgb_rays.append(rgb)\n",
    "                depth_rays_preds.append(depth_pred)\n",
    "\n",
    "            \n",
    "            depth_rays_preds = np.concatenate(depth_rays_preds).reshape(H, W)\n",
    "            depth_rays_preds, _ = visualize_depth_numpy(depth_rays_preds, near_far_source)\n",
    "            \n",
    "            rgb_rays = np.concatenate(rgb_rays).reshape(H, W, 3)\n",
    "            img_vis = np.concatenate((img*255,rgb_rays*255,depth_rays_preds),axis=1)\n",
    "            \n",
    "            if save_as_image:\n",
    "                imageio.imwrite(f'{save_dir}/{scene}_{val_idx[i]:03d}.png', img_vis.astype('uint8'))\n",
    "            else:\n",
    "                rgbs.append(img_vis.astype('uint8'))\n",
    "                \n",
    "            # quantity\n",
    "            # center crop 0.8 ratio\n",
    "            H_crop, W_crop = np.array(rgb_rays.shape[:2])//10\n",
    "            img = img[H_crop:-H_crop,W_crop:-W_crop]\n",
    "            rgb_rays = rgb_rays[H_crop:-H_crop,W_crop:-W_crop]\n",
    "            \n",
    "            psnr.append( mse2psnr(np.mean((rgb_rays-img)**2)))\n",
    "            ssim.append( structural_similarity(rgb_rays, img, multichannel=True))\n",
    "            \n",
    "            img_tensor = torch.from_numpy(rgb_rays)[None].permute(0,3,1,2).float()*2-1.0 # image should be RGB, IMPORTANT: normalized to [-1,1]\n",
    "            img_gt_tensor = torch.from_numpy(img)[None].permute(0,3,1,2).float()*2-1.0\n",
    "            LPIPS_vgg.append( loss_fn_vgg(img_tensor, img_gt_tensor).item())\n",
    "        \n",
    "        print(f'=====> scene: {scene} mean psnr {np.mean(psnr)} ssim: {np.mean(ssim)} lpips: {np.mean(LPIPS_vgg)}')   \n",
    "        psnr_all.append(psnr);ssim_all.append(ssim);LPIPS_vgg_all.append(LPIPS_vgg)\n",
    "        \n",
    "    if not save_as_image:\n",
    "        imageio.mimwrite(f'{save_dir}/{scene}_spiral.mp4', np.stack(rgbs), fps=20, quality=10)\n",
    "print(f'=====> all mean psnr {np.mean(psnr_all)} ssim: {np.mean(ssim_all)} lpips: {np.mean(LPIPS_vgg_all)}') "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# nerf no fine tuning"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## rendering novel views with nearest 3 views"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "psnr_all,ssim_all,LPIPS_vgg_all = [],[],[]\n",
    "for i_scene, scene in enumerate(['ship','mic','chair','lego','drums','ficus','materials','hotdog']):#\n",
    "    psnr,ssim,LPIPS_vgg = [],[],[]\n",
    "    cmd = f'--datadir /mnt/new_disk_2/anpei/Dataset/nerf_synthetic/{scene}  \\\n",
    "     --dataset_name blender --white_bkgd \\\n",
    "    --net_type v0 --ckpt ./ckpts/mvsnerf-v0.tar '\n",
    "\n",
    "    args = config_parser(cmd.split())\n",
    "    args.use_viewdirs = True\n",
    "\n",
    "    args.N_samples = 128\n",
    "    args.feat_dim =  8+12\n",
    "\n",
    "    # create models\n",
    "    if 0==i_scene:\n",
    "        render_kwargs_train, render_kwargs_test, start, grad_vars = create_nerf_mvs(args, use_mvs=True, dir_embedder=False, pts_embedder=True)\n",
    "        filter_keys(render_kwargs_train)\n",
    "\n",
    "        MVSNet = render_kwargs_train['network_mvs']\n",
    "        render_kwargs_train.pop('network_mvs')\n",
    "\n",
    "\n",
    "    datadir = args.datadir\n",
    "    datatype = 'train'\n",
    "    pad = 16\n",
    "    args.chunk = 5120\n",
    "\n",
    "\n",
    "    print('============> rendering dataset <===================')\n",
    "    dataset_train = dataset_dict[args.dataset_name](args, split='train')\n",
    "    dataset_val = dataset_dict[args.dataset_name](args, split='val')\n",
    "    val_idx = dataset_val.img_idx\n",
    "    \n",
    "    save_as_image = True\n",
    "    save_dir = f'results/test3'\n",
    "    os.makedirs(save_dir, exist_ok=True)\n",
    "    MVSNet.train()\n",
    "    MVSNet = MVSNet.cuda()\n",
    "    \n",
    "    with torch.no_grad():\n",
    "\n",
    "        try:\n",
    "            tqdm._instances.clear() \n",
    "        except Exception:     \n",
    "            pass\n",
    "        \n",
    "        for i, batch in enumerate(tqdm(dataset_val)):\n",
    "            torch.cuda.empty_cache()\n",
    "\n",
    "            \n",
    "            rays, img = decode_batch(batch)\n",
    "            rays = rays.squeeze().to(device)  # (H*W, 3)\n",
    "            img = img.squeeze().cpu().numpy()  # (H, W, 3)\n",
    "        \n",
    "            # find nearest image idx from training views\n",
    "            positions = dataset_train.poses[:,:3,3]\n",
    "            dis = np.sum(np.abs(positions - dataset_val.poses[[i],:3,3]), axis=-1)\n",
    "            pair_idx = np.argsort(dis)[:3]\n",
    "            pair_idx = [dataset_train.img_idx[item] for item in pair_idx]\n",
    "            \n",
    "            imgs_source, proj_mats, near_far_source, pose_source = dataset_train.read_source_views(pair_idx=pair_idx,device=device)\n",
    "            volume_feature, _, _ = MVSNet(imgs_source, proj_mats, near_far_source, pad=pad)\n",
    "            imgs_source = unpreprocess(imgs_source)\n",
    "        \n",
    "            N_rays_all = rays.shape[0]\n",
    "            rgb_rays, depth_rays_preds = [],[]\n",
    "            for chunk_idx in range(N_rays_all//args.chunk + int(N_rays_all%args.chunk>0)):\n",
    "\n",
    "                xyz_coarse_sampled, rays_o, rays_d, z_vals = ray_marcher(rays[chunk_idx*args.chunk:(chunk_idx+1)*args.chunk],\n",
    "                                                    N_samples=args.N_samples)\n",
    "\n",
    "                # Converting world coordinate to ndc coordinate\n",
    "                H, W = img.shape[:2]\n",
    "                inv_scale = torch.tensor([W - 1, H - 1]).to(device)\n",
    "                w2c_ref, intrinsic_ref = pose_source['w2cs'][0], pose_source['intrinsics'][0].clone()\n",
    "                intrinsic_ref[:2] *= args.imgScale_test/args.imgScale_train\n",
    "                xyz_NDC = get_ndc_coordinate(w2c_ref, intrinsic_ref, xyz_coarse_sampled, inv_scale,\n",
    "                                             near=near_far_source[0], far=near_far_source[1], pad=pad*args.imgScale_test)\n",
    "\n",
    "\n",
    "                # rendering\n",
    "                rgb, disp, acc, depth_pred, alpha, extras = rendering(args, pose_source, xyz_coarse_sampled,\n",
    "                                                                       xyz_NDC, z_vals, rays_o, rays_d,\n",
    "                                                                       volume_feature,imgs_source, **render_kwargs_train)\n",
    "    \n",
    "                \n",
    "                rgb, depth_pred = torch.clamp(rgb.cpu(),0,1.0).numpy(), depth_pred.cpu().numpy()\n",
    "                rgb_rays.append(rgb)\n",
    "                depth_rays_preds.append(depth_pred)\n",
    "\n",
    "            \n",
    "            depth_rays_preds = np.concatenate(depth_rays_preds).reshape(H, W)\n",
    "            depth_rays_preds, _ = visualize_depth_numpy(depth_rays_preds, near_far_source)\n",
    "            \n",
    "            rgb_rays = np.concatenate(rgb_rays).reshape(H, W, 3)\n",
    "            img_vis = np.concatenate((img*255,rgb_rays*255,depth_rays_preds),axis=1)\n",
    "            \n",
    "#             img_vis = np.concatenate((torch.cat(torch.split(imgs_source*255, [1,1,1], dim=1),-1).squeeze().permute(1,2,0).cpu().numpy(),img_vis),axis=1)\n",
    "            \n",
    "            if save_as_image:\n",
    "                imageio.imwrite(f'{save_dir}/{scene}_{val_idx[i]:03d}.png', img_vis.astype('uint8'))\n",
    "            else:\n",
    "                rgbs.append(img_vis.astype('uint8'))\n",
    "                \n",
    "            # quantity\n",
    "            # center crop 0.8 ratio\n",
    "            H_crop, W_crop = np.array(rgb_rays.shape[:2])//10\n",
    "            img = img[H_crop:-H_crop,W_crop:-W_crop]\n",
    "            rgb_rays = rgb_rays[H_crop:-H_crop,W_crop:-W_crop]\n",
    "            \n",
    "            psnr.append( mse2psnr(np.mean((rgb_rays-img)**2)))\n",
    "            ssim.append( structural_similarity(rgb_rays, img, multichannel=True))\n",
    "            \n",
    "            img_tensor = torch.from_numpy(rgb_rays)[None].permute(0,3,1,2).float()*2-1.0 # image should be RGB, IMPORTANT: normalized to [-1,1]\n",
    "            img_gt_tensor = torch.from_numpy(img)[None].permute(0,3,1,2).float()*2-1.0\n",
    "            LPIPS_vgg.append( loss_fn_vgg(img_tensor, img_gt_tensor).item())\n",
    "\n",
    "        print(f'=====> scene: {scene} mean psnr {np.mean(psnr)} ssim: {np.mean(ssim)} lpips: {np.mean(LPIPS_vgg)}')   \n",
    "        psnr_all.append(psnr);ssim_all.append(ssim);LPIPS_vgg_all.append(LPIPS_vgg)\n",
    "\n",
    "    if not save_as_image:\n",
    "        imageio.mimwrite(f'{save_dir}/{scene}_spiral.mp4', np.stack(rgbs), fps=20, quality=10)\n",
    "\n",
    "print(f'=====> all mean psnr {np.mean(psnr_all)} ssim: {np.mean(ssim_all)} lpips: {np.mean(LPIPS_vgg_all)}') "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## rendering novel views with fixed 3 source views"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "psnr_all,ssim_all,LPIPS_vgg_all = [],[],[]\n",
    "for i_scene, scene in enumerate(['ship','mic','chair','lego','drums','ficus','materials','hotdog']):#'ship','mic','chair','lego','drums','ficus','materials','hotdog'\n",
    "    psnr,ssim,LPIPS_vgg = [],[],[]\n",
    "    cmd = f'--datadir /mnt/new_disk_2/anpei/Dataset/nerf_synthetic/{scene}  \\\n",
    "     --dataset_name blender --white_bkgd \\\n",
    "     --ckpt /mnt/new_disk_2/anpei/code/MVS-NeRF/runs_new/mvs-nerf-color-skip-no-border/ckpts/79999.tar \\\n",
    "     --net_type v2 --netwidth 128 --netdepth 6'\n",
    "\n",
    "    args = config_parser(cmd.split())\n",
    "    args.use_viewdirs = True\n",
    "\n",
    "    args.N_samples = 128\n",
    "    args.feat_dim = 8+3*4\n",
    "\n",
    "    # create models\n",
    "    if 0==i_scene:\n",
    "        render_kwargs_train, render_kwargs_test, start, grad_vars = create_nerf_mvs(args, use_mvs=True, dir_embedder=False, pts_embedder=True)\n",
    "        filter_keys(render_kwargs_train)\n",
    "\n",
    "        MVSNet = render_kwargs_train['network_mvs']\n",
    "        render_kwargs_train.pop('network_mvs')\n",
    "\n",
    "\n",
    "    datadir = args.datadir\n",
    "    datatype = 'val'\n",
    "    pad = 16\n",
    "    args.chunk = 5120\n",
    "\n",
    "\n",
    "    print('============> rendering dataset <===================')\n",
    "    dataset = dataset_dict[args.dataset_name](args, split=datatype)\n",
    "    val_idx = dataset.img_idx\n",
    "    \n",
    "    save_as_image = True\n",
    "    save_dir = f'results/test3'\n",
    "    os.makedirs(save_dir, exist_ok=True)\n",
    "    MVSNet.train()\n",
    "    MVSNet = MVSNet.cuda()\n",
    "    \n",
    "    with torch.no_grad():\n",
    "\n",
    "        imgs_source, proj_mats, near_far_source, pose_source = dataset.read_source_views(device=device)\n",
    "        volume_feature, _, _ = MVSNet(imgs_source, proj_mats, near_far_source, pad=pad)\n",
    "        imgs_source = unpreprocess(imgs_source)\n",
    "\n",
    "        try:\n",
    "            tqdm._instances.clear() \n",
    "        except Exception:     \n",
    "            pass\n",
    "        \n",
    "        for i, batch in enumerate(tqdm(dataset)):\n",
    "            torch.cuda.empty_cache()\n",
    "            \n",
    "            rays, img = decode_batch(batch)\n",
    "            rays = rays.squeeze().to(device)  # (H*W, 3)\n",
    "            img = img.squeeze().cpu().numpy()  # (H, W, 3)\n",
    "        \n",
    "            N_rays_all = rays.shape[0]\n",
    "            rgb_rays, depth_rays_preds = [],[]\n",
    "            for chunk_idx in range(N_rays_all//args.chunk + int(N_rays_all%args.chunk>0)):\n",
    "\n",
    "                xyz_coarse_sampled, rays_o, rays_d, z_vals = ray_marcher(rays[chunk_idx*args.chunk:(chunk_idx+1)*args.chunk],\n",
    "                                                    N_samples=args.N_samples)\n",
    "\n",
    "                # Converting world coordinate to ndc coordinate\n",
    "                H, W = img.shape[:2]\n",
    "                inv_scale = torch.tensor([W - 1, H - 1]).to(device)\n",
    "                w2c_ref, intrinsic_ref = pose_source['w2cs'][0], pose_source['intrinsics'][0].clone()\n",
    "                xyz_NDC = get_ndc_coordinate(w2c_ref, intrinsic_ref, xyz_coarse_sampled, inv_scale,\n",
    "                                             near=near_far_source[0], far=near_far_source[1], pad=pad*args.imgScale_test)\n",
    "\n",
    "\n",
    "                # rendering\n",
    "                rgb, disp, acc, depth_pred, alpha, extras = rendering(args, pose_source, xyz_coarse_sampled,\n",
    "                                                                       xyz_NDC, z_vals, rays_o, rays_d,\n",
    "                                                                       volume_feature,imgs_source, **render_kwargs_train)\n",
    "    \n",
    "                \n",
    "                rgb, depth_pred = torch.clamp(rgb.cpu(),0,1.0).numpy(), depth_pred.cpu().numpy()\n",
    "                rgb_rays.append(rgb)\n",
    "                depth_rays_preds.append(depth_pred)\n",
    "\n",
    "            \n",
    "            depth_rays_preds = np.concatenate(depth_rays_preds).reshape(H, W)\n",
    "            depth_rays_preds, _ = visualize_depth_numpy(depth_rays_preds, near_far_source)\n",
    "            \n",
    "            rgb_rays = np.concatenate(rgb_rays).reshape(H, W, 3)\n",
    "            img_vis = np.concatenate((img*255,rgb_rays*255,depth_rays_preds),axis=1)\n",
    "            \n",
    "            img_vis = np.concatenate((torch.cat(torch.split(imgs_source*255, [1,1,1], dim=1),-1).squeeze().permute(1,2,0).cpu().numpy(),img_vis),axis=1)\n",
    "            \n",
    "            if save_as_image:\n",
    "                imageio.imwrite(f'{save_dir}/{scene}_{val_idx[i]:03d}.png', img_vis.astype('uint8'))\n",
    "            else:\n",
    "                rgbs.append(img_vis.astype('uint8'))\n",
    "                \n",
    "            # quantity\n",
    "            psnr.append( mse2psnr(np.mean((rgb_rays-img)**2)))\n",
    "            ssim.append( structural_similarity(rgb_rays, img, multichannel=True))\n",
    "            \n",
    "            img_tensor = torch.from_numpy(rgb_rays)[None].permute(0,3,1,2).float()*2-1.0 # image should be RGB, IMPORTANT: normalized to [-1,1]\n",
    "            img_gt_tensor = torch.from_numpy(img)[None].permute(0,3,1,2).float()*2-1.0\n",
    "            LPIPS_vgg.append( loss_fn_vgg(img_tensor, img_gt_tensor).item())\n",
    "\n",
    "        print(f'=====> scene: {scene} mean psnr {np.mean(psnr)} ssim: {np.mean(ssim)} lpips: {np.mean(LPIPS_vgg)}')   \n",
    "        psnr_all.append(psnr);ssim_all.append(ssim);LPIPS_vgg_all.append(LPIPS_vgg)\n",
    "\n",
    "    if not save_as_image:\n",
    "        imageio.mimwrite(f'{save_dir}/{scene}_depth_spiral.mp4', np.stack(depths_vis), fps=10, quality=10)\n",
    "        imageio.mimwrite(f'{save_dir}/{scene}_spiral.mp4', np.stack(rgbs), fps=20, quality=10)\n",
    "print(f'=====> all mean psnr {np.mean(psnr_all)} ssim: {np.mean(ssim_all)} lpips: {np.mean(LPIPS_vgg_all)}') "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# DTU no fine tuning"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## rendering novel views with nearest 3 views"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "psnr_all,ssim_all,LPIPS_vgg_all = [],[],[]\n",
    "depth_acc = {}\n",
    "eval_metric = [0.1,0.05,0.01]\n",
    "depth_acc[f'abs_err'],depth_acc[f'acc_l_{eval_metric[0]}'],depth_acc[f'acc_l_{eval_metric[1]}'],depth_acc[f'acc_l_{eval_metric[2]}'] = {},{},{},{}\n",
    "for i_scene, scene in enumerate([1,8,21,103,114]):#,8,21,103,114\n",
    "    psnr,ssim,LPIPS_vgg = [],[],[]\n",
    "    cmd = f'--datadir /mnt/data/new_disk/sungx/data/mvs_dataset/DTU/mvs_training/dtu/scan{scene}  \\\n",
    "     --dataset_name dtu_ft  \\\n",
    "     --net_type v0 --ckpt ./ckpts//mvsnerf-v0.tar '\n",
    "\n",
    "    args = config_parser(cmd.split())\n",
    "    args.use_viewdirs = True\n",
    "\n",
    "    args.N_samples = 128\n",
    "    args.feat_dim =  8+12\n",
    "\n",
    "    # create models\n",
    "    if 0==i_scene:\n",
    "        render_kwargs_train, render_kwargs_test, start, grad_vars = create_nerf_mvs(args, use_mvs=True, dir_embedder=False, pts_embedder=True)\n",
    "        filter_keys(render_kwargs_train)\n",
    "\n",
    "        MVSNet = render_kwargs_train['network_mvs']\n",
    "        render_kwargs_train.pop('network_mvs')\n",
    "\n",
    "\n",
    "    datadir = args.datadir\n",
    "    datatype = 'train'\n",
    "    pad = 16\n",
    "    args.chunk = 5120\n",
    "\n",
    "\n",
    "    print('============> rendering dataset <===================')\n",
    "    dataset_train = dataset_dict[args.dataset_name](args, split='train')\n",
    "    dataset_val = dataset_dict[args.dataset_name](args, split='val')\n",
    "    val_idx = dataset_val.img_idx\n",
    "    \n",
    "    save_as_image = True\n",
    "    save_dir = f'results/test3'\n",
    "    os.makedirs(save_dir, exist_ok=True)\n",
    "    MVSNet.train()\n",
    "    MVSNet = MVSNet.cuda()\n",
    "    \n",
    "\n",
    "\n",
    "    with torch.no_grad():\n",
    "\n",
    "        try:\n",
    "            tqdm._instances.clear() \n",
    "        except Exception:     \n",
    "            pass\n",
    "        \n",
    "        for i, batch in enumerate(tqdm(dataset_val)):\n",
    "            torch.cuda.empty_cache()\n",
    "            \n",
    "            rays, img = decode_batch(batch)\n",
    "            rays = rays.squeeze().to(device)  # (H*W, 3)\n",
    "            img = img.squeeze().cpu().numpy()  # (H, W, 3)\n",
    "            depth = batch['depth'].squeeze().numpy()  # (H, W)\n",
    "        \n",
    "            # find nearest image idx from training views\n",
    "            positions = dataset_train.poses[:,:3,3]\n",
    "            dis = np.sum(np.abs(positions - dataset_val.poses[[i],:3,3]), axis=-1)\n",
    "            pair_idx = np.argsort(dis)[:3]\n",
    "            pair_idx = [dataset_train.img_idx[item] for item in pair_idx]\n",
    "            \n",
    "            imgs_source, proj_mats, near_far_source, pose_source = dataset_train.read_source_views(pair_idx=pair_idx,device=device)\n",
    "            volume_feature, _, _ = MVSNet(imgs_source, proj_mats, near_far_source, pad=pad)\n",
    "            imgs_source = unpreprocess(imgs_source)\n",
    "        \n",
    "            N_rays_all = rays.shape[0]\n",
    "            rgb_rays, depth_rays_preds = [],[]\n",
    "            for chunk_idx in range(N_rays_all//args.chunk + int(N_rays_all%args.chunk>0)):\n",
    "\n",
    "                xyz_coarse_sampled, rays_o, rays_d, z_vals = ray_marcher(rays[chunk_idx*args.chunk:(chunk_idx+1)*args.chunk],\n",
    "                                                    N_samples=args.N_samples)\n",
    "\n",
    "                # Converting world coordinate to ndc coordinate\n",
    "                H, W = img.shape[:2]\n",
    "                inv_scale = torch.tensor([W - 1, H - 1]).to(device)\n",
    "                w2c_ref, intrinsic_ref = pose_source['w2cs'][0], pose_source['intrinsics'][0].clone()\n",
    "                xyz_NDC = get_ndc_coordinate(w2c_ref, intrinsic_ref, xyz_coarse_sampled, inv_scale,\n",
    "                                             near=near_far_source[0], far=near_far_source[1], pad=pad*args.imgScale_test)\n",
    "\n",
    "\n",
    "                # rendering\n",
    "                rgb, disp, acc, depth_pred, alpha, extras = rendering(args, pose_source, xyz_coarse_sampled,\n",
    "                                                                       xyz_NDC, z_vals, rays_o, rays_d,\n",
    "                                                                       volume_feature,imgs_source, **render_kwargs_train)\n",
    "    \n",
    "                \n",
    "                rgb, depth_pred = torch.clamp(rgb.cpu(),0,1.0).numpy(), depth_pred.cpu().numpy()\n",
    "                rgb_rays.append(rgb)\n",
    "                depth_rays_preds.append(depth_pred)\n",
    "\n",
    "            \n",
    "            depth_rays_preds = np.concatenate(depth_rays_preds).reshape(H, W)\n",
    "\n",
    "            depth_gt, _ =  read_depth(f'/mnt/data/new_disk/sungx/data/mvs_dataset/DTU/mvs_training/dtu/Depths/scan{scene}/depth_map_{val_idx[i]:04d}.pfm')\n",
    "        \n",
    "            mask_gt = depth_gt>0\n",
    "            abs_err = abs_error(depth_rays_preds, depth_gt/200, mask_gt)\n",
    "\n",
    "            eval_metric = [0.01,0.05, 0.1]\n",
    "            depth_acc[f'abs_err'][f'{scene}'] = np.mean(abs_err)\n",
    "            depth_acc[f'acc_l_{eval_metric[0]}'][f'{scene}'] = acc_threshold(abs_err,eval_metric[0]).mean()\n",
    "            depth_acc[f'acc_l_{eval_metric[1]}'][f'{scene}'] = acc_threshold(abs_err,eval_metric[1]).mean()\n",
    "            depth_acc[f'acc_l_{eval_metric[2]}'][f'{scene}'] = acc_threshold(abs_err,eval_metric[2]).mean()\n",
    "\n",
    "            \n",
    "            depth_rays_preds, _ = visualize_depth_numpy(depth_rays_preds, near_far_source)\n",
    "            \n",
    "            rgb_rays = np.concatenate(rgb_rays).reshape(H, W, 3)\n",
    "            img_vis = np.concatenate((img*255,rgb_rays*255,depth_rays_preds),axis=1)\n",
    "            \n",
    "            if save_as_image:\n",
    "                imageio.imwrite(f'{save_dir}/scan{scene}_{val_idx[i]:03d}.png', img_vis.astype('uint8'))\n",
    "            else:\n",
    "                rgbs.append(img_vis.astype('uint8'))\n",
    "                \n",
    "            # quantity\n",
    "            # mask background since they are outside the far boundle\n",
    "            mask = depth==0\n",
    "            imageio.imwrite(f'{save_dir}/scan{scene}_{val_idx[i]:03d}_mask.png', mask.astype('uint8')*255)\n",
    "            rgb_rays[mask],img[mask] = 0.0,0.0\n",
    "            psnr.append( mse2psnr(np.mean((rgb_rays[~mask]-img[~mask])**2)))\n",
    "            ssim.append( structural_similarity(rgb_rays, img, multichannel=True))\n",
    "            \n",
    "            img_tensor = torch.from_numpy(rgb_rays)[None].permute(0,3,1,2).float()*2-1.0 # image should be RGB, IMPORTANT: normalized to [-1,1]\n",
    "            img_gt_tensor = torch.from_numpy(img)[None].permute(0,3,1,2).float()*2-1.0\n",
    "            LPIPS_vgg.append( loss_fn_vgg(img_tensor, img_gt_tensor).item())\n",
    "\n",
    "        print(f'=====> scene: {scene} mean psnr {np.mean(psnr)} ssim: {np.mean(ssim)} lpips: {np.mean(LPIPS_vgg)}')   \n",
    "        psnr_all.append(psnr);ssim_all.append(ssim);LPIPS_vgg_all.append(LPIPS_vgg)\n",
    "        \n",
    "\n",
    "    if not save_as_image:\n",
    "        imageio.mimwrite(f'{save_dir}/{scene}_spiral.mp4', np.stack(rgbs), fps=20, quality=10)\n",
    "\n",
    "a = np.mean(list(depth_acc['abs_err'].values()))\n",
    "b = np.mean(list(depth_acc[f'acc_l_{eval_metric[0]}'].values()))\n",
    "c = np.mean(list(depth_acc[f'acc_l_{eval_metric[1]}'].values()))\n",
    "d = np.mean(list(depth_acc[f'acc_l_{eval_metric[2]}'].values()))\n",
    "print(f'============> abs_err: {a} <=================')\n",
    "print(f'============> acc_l_{eval_metric[0]}: {b} <=================')\n",
    "print(f'============> acc_l_{eval_metric[1]}: {c} <=================')\n",
    "print(f'============> acc_l_{eval_metric[2]}: {d} <=================')\n",
    "print(f'=====> all mean psnr {np.mean(psnr_all)} ssim: {np.mean(ssim_all)} lpips: {np.mean(LPIPS_vgg_all)}') "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## rendering novel views with fixed 3 source views"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "psnr_all,ssim_all,LPIPS_vgg_all = [],[],[]\n",
    "for i_scene, scene in enumerate([1,8,21,103,114]):#,\n",
    "    psnr,ssim,LPIPS_vgg = [],[],[]\n",
    "    cmd = f'--datadir /mnt/data/new_disk/sungx/data/mvs_dataset/DTU/mvs_training/dtu/scan{scene}  \\\n",
    "    --dataset_name dtu_ft  \\\n",
    "    --ckpt ./ckpts//mvsnerf-v0.tar'\n",
    "\n",
    "    args = config_parser(cmd.split())\n",
    "    args.use_viewdirs = True\n",
    "\n",
    "    args.N_samples = 128\n",
    "    args.feat_dim =  8+3*4\n",
    "\n",
    "    # create models\n",
    "    if 0==i_scene:\n",
    "        render_kwargs_train, render_kwargs_test, start, grad_vars = create_nerf_mvs(args, use_mvs=True, dir_embedder=False, pts_embedder=True)\n",
    "        filter_keys(render_kwargs_train)\n",
    "\n",
    "        MVSNet = render_kwargs_train['network_mvs']\n",
    "        render_kwargs_train.pop('network_mvs')\n",
    "\n",
    "\n",
    "    datadir = args.datadir\n",
    "    datatype = 'val'\n",
    "    pad = 24\n",
    "    args.chunk = 5120\n",
    "\n",
    "\n",
    "    print('============> rendering dataset <===================')\n",
    "    dataset = dataset_dict[args.dataset_name](args, split=datatype)\n",
    "    val_idx = dataset.img_idx\n",
    "    \n",
    "    save_as_image = True\n",
    "    save_dir = f'results/test3'\n",
    "    os.makedirs(save_dir, exist_ok=True)\n",
    "    MVSNet.train()\n",
    "    MVSNet = MVSNet.cuda()\n",
    "    \n",
    "    with torch.no_grad():\n",
    "\n",
    "        imgs_source, proj_mats, near_far_source, pose_source = dataset.read_source_views(device=device)\n",
    "        volume_feature, _, _ = MVSNet(imgs_source, proj_mats, near_far_source, pad=pad)\n",
    "        imgs_source = unpreprocess(imgs_source)\n",
    "\n",
    "        try:\n",
    "            tqdm._instances.clear() \n",
    "        except Exception:     \n",
    "            pass\n",
    "        \n",
    "        for i, batch in enumerate(tqdm(dataset)):\n",
    "            torch.cuda.empty_cache()\n",
    "            \n",
    "            rays, img = decode_batch(batch)\n",
    "            rays = rays.squeeze().to(device)  # (H*W, 3)\n",
    "            img = img.squeeze().cpu().numpy()  # (H, W, 3)\n",
    "            depth = batch['depth'].squeeze().numpy()  # (H, W)\n",
    "        \n",
    "            N_rays_all = rays.shape[0]\n",
    "            rgb_rays, depth_rays_preds = [],[]\n",
    "            for chunk_idx in range(N_rays_all//args.chunk + int(N_rays_all%args.chunk>0)):\n",
    "\n",
    "                xyz_coarse_sampled, rays_o, rays_d, z_vals = ray_marcher(rays[chunk_idx*args.chunk:(chunk_idx+1)*args.chunk],\n",
    "                                                    N_samples=args.N_samples)\n",
    "\n",
    "                # Converting world coordinate to ndc coordinate\n",
    "                H, W = img.shape[:2]\n",
    "                inv_scale = torch.tensor([W - 1, H - 1]).to(device)\n",
    "                w2c_ref, intrinsic_ref = pose_source['w2cs'][0], pose_source['intrinsics'][0].clone()\n",
    "                xyz_NDC = get_ndc_coordinate(w2c_ref, intrinsic_ref, xyz_coarse_sampled, inv_scale,\n",
    "                                             near=near_far_source[0], far=near_far_source[1], pad=pad*args.imgScale_test)\n",
    "\n",
    "\n",
    "                # rendering\n",
    "                rgb, disp, acc, depth_pred, alpha, extras = rendering(args, pose_source, xyz_coarse_sampled,\n",
    "                                                                       xyz_NDC, z_vals, rays_o, rays_d,\n",
    "                                                                       volume_feature,imgs_source, **render_kwargs_train)\n",
    "    \n",
    "                \n",
    "                rgb, depth_pred = torch.clamp(rgb.cpu(),0,1.0).numpy(), depth_pred.cpu().numpy()\n",
    "                rgb_rays.append(rgb)\n",
    "                depth_rays_preds.append(depth_pred)\n",
    "\n",
    "            \n",
    "            depth_rays_preds = np.concatenate(depth_rays_preds).reshape(H, W)\n",
    "            depth_rays_preds, _ = visualize_depth_numpy(depth_rays_preds, near_far_source)\n",
    "            \n",
    "            rgb_rays = np.concatenate(rgb_rays).reshape(H, W, 3)\n",
    "            img_vis = np.concatenate((img*255,rgb_rays*255,depth_rays_preds),axis=1)\n",
    "            \n",
    "            if save_as_image:\n",
    "                imageio.imwrite(f'{save_dir}/scan{scene}_{val_idx[i]:03d}.png', img_vis.astype('uint8'))\n",
    "            else:\n",
    "                rgbs.append(img_vis.astype('uint8'))\n",
    "                \n",
    "            # quantity\n",
    "            # mask background since they are outside the far boundle\n",
    "            mask = depth==0\n",
    "            imageio.imwrite(f'{save_dir}/scan{scene}_{val_idx[i]:03d}_mask.png', mask.astype('uint8')*255)\n",
    "            rgb_rays[mask],img[mask] = 0.0,0.0\n",
    "            psnr.append( mse2psnr(np.mean((rgb_rays[~mask]-img[~mask])**2)))\n",
    "            ssim.append( structural_similarity(rgb_rays, img, multichannel=True))\n",
    "            \n",
    "            img_tensor = torch.from_numpy(rgb_rays)[None].permute(0,3,1,2).float()*2-1.0 # image should be RGB, IMPORTANT: normalized to [-1,1]\n",
    "            img_gt_tensor = torch.from_numpy(img)[None].permute(0,3,1,2).float()*2-1.0\n",
    "            LPIPS_vgg.append( loss_fn_vgg(img_tensor, img_gt_tensor).item())\n",
    "\n",
    "        print(f'=====> scene: {scene} mean psnr {np.mean(psnr)} ssim: {np.mean(ssim)} lpips: {np.mean(LPIPS_vgg)}')   \n",
    "        psnr_all.append(psnr);ssim_all.append(ssim);LPIPS_vgg_all.append(LPIPS_vgg)\n",
    "\n",
    "    if not save_as_image:\n",
    "        imageio.mimwrite(f'{save_dir}/{scene}_spiral.mp4', np.stack(rgbs), fps=20, quality=10)\n",
    "print(f'=====> all mean psnr {np.mean(psnr_all)} ssim: {np.mean(ssim_all)} lpips: {np.mean(LPIPS_vgg_all)}') "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Pairs generation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json,torch\n",
    "import sys,os\n",
    "import numpy as np\n",
    "root = '/home/hengfei/Desktop/research/mvsnerf'\n",
    "os.chdir(root)\n",
    "sys.path.append(root)\n",
    "pairs = torch.load('./configs/pairs.th')\n",
    "\n",
    "# llff\n",
    "root_dir = '/home/hengfei/Desktop/research/mvsnerf/xgaze/'\n",
    "for scene in ['xgaze_11images_cropped_colmapCODE']:#\n",
    "    poses_bounds = np.load(os.path.join(root_dir, scene, 'poses_bounds.npy'))  # (N_images, 11)\n",
    "    poses = poses_bounds[:, :15].reshape(-1, 3, 5)  # (N_images, 3, 5)\n",
    "    poses = np.concatenate([poses[..., 1:2], - poses[..., :1], poses[..., 2:4]], -1)\n",
    "\n",
    "    ref_position = np.mean(poses[..., 3],axis=0, keepdims=True)\n",
    "    dist = np.sum(np.abs(poses[..., 3] - ref_position), axis=-1)\n",
    "    pair_idx = np.argsort(dist)[:11]\n",
    "#     pair_idx = torch.randperm(len(poses))[:20].tolist()\n",
    "\n",
    "    pairs[f'{scene}_test'] = pair_idx[::6]\n",
    "    pairs[f'{scene}_val'] = pair_idx[::6]\n",
    "    pairs[f'{scene}_train'] = np.delete(pair_idx, range(0,11,6))\n",
    "\n",
    "torch.save(pairs,'/home/hengfei/Desktop/research/mvsnerf/configs/pairs.th')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# quantity evauation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys,os,imageio,lpips,cv2,torch,glob\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from skimage.metrics import structural_similarity"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Setting up [LPIPS] perceptual loss: trunk [vgg], v[0.1], spatial [off]\n",
      "Loading model from: /home/anpei/anaconda3/lib/python3.7/site-packages/lpips/weights/v0.1/vgg.pth\n"
     ]
    }
   ],
   "source": [
    "loss_fn_vgg = lpips.LPIPS(net='vgg') \n",
    "mse2psnr = lambda x : -10. * np.log(x) / np.log(10.)\n",
    "\n",
    "\n",
    "\n",
    "def acc_threshold(abs_err, threshold):\n",
    "    \"\"\"\n",
    "    computes the percentage of pixels whose depth error is less than @threshold\n",
    "    \"\"\"\n",
    "    acc_mask = abs_err < threshold\n",
    "    return  acc_mask.astype('float') if type(abs_err) is np.ndarray else acc_mask.float()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# nerf "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "=====> scene: chair mean psnr 31.070870959993407 ssim: 0.970869913728838 lpips: 0.05512496456503868\n",
      "=====> scene: drums mean psnr 25.464383523724557 ssim: 0.9430287321705997 lpips: 0.1010842639952898\n",
      "=====> scene: ficus mean psnr 29.72717081186501 ssim: 0.9688198661712594 lpips: 0.04721927270293236\n",
      "=====> scene: hotdog mean psnr 34.63162021512352 ssim: 0.9798700143526428 lpips: 0.0885334312915802\n",
      "=====> scene: lego mean psnr 32.65761069614622 ssim: 0.9751430050524844 lpips: 0.05375238787382841\n",
      "=====> scene: materials mean psnr 30.220202654922936 ssim: 0.9677394226502894 lpips: 0.1052329633384943\n",
      "=====> scene: mic mean psnr 31.810551677509977 ssim: 0.9810118386928188 lpips: 0.03268271638080478\n",
      "=====> scene: ship mean psnr 29.487980342358682 ssim: 0.9079920156014059 lpips: 0.2625834122300148\n",
      "=====> all mean psnr 30.633798860205538 ssim: 0.9618093510525423 lpips: 0.09327667654724792\n",
      "=====> scene: fern mean psnr 23.87081932481545 ssim: 0.828319405500272 lpips: 0.29106350988149643\n",
      "=====> scene: flower mean psnr 26.84248375485232 ssim: 0.8972176342834637 lpips: 0.175886869430542\n",
      "=====> scene: fortress mean psnr 31.368287455491632 ssim: 0.9454387223441942 lpips: 0.14687807857990265\n",
      "=====> scene: horns mean psnr 25.957297279910254 ssim: 0.9002932801372645 lpips: 0.24746065214276314\n",
      "=====> scene: leaves mean psnr 21.209230306630403 ssim: 0.7924383925048237 lpips: 0.3013022392988205\n",
      "=====> scene: orchids mean psnr 19.805083676014405 ssim: 0.7216119459229409 lpips: 0.3210122361779213\n",
      "=====> scene: room mean psnr 33.538893189163886 ssim: 0.9782926576909046 lpips: 0.15664737671613693\n",
      "=====> scene: trex mean psnr 25.191124840989897 ssim: 0.8986906362165991 lpips: 0.24522138014435768\n",
      "=====> all mean psnr 25.972902478483533 ssim: 0.8702878343250579 lpips: 0.23568404279649258\n",
      "=====> scene: 1 mean psnr 26.621467948629075 ssim: 0.9015109955368615 lpips: 0.26542308926582336\n",
      "=====> scene: 8 mean psnr 28.331380911917446 ssim: 0.8758834032089559 lpips: 0.32112571597099304\n",
      "=====> scene: 21 mean psnr 23.238617049537122 ssim: 0.8736486067576469 lpips: 0.2456555962562561\n",
      "=====> scene: 103 mean psnr 30.40125554666663 ssim: 0.9442875531229784 lpips: 0.25619567558169365\n",
      "=====> scene: 114 mean psnr 26.46890004323232 ssim: 0.9134870027630088 lpips: 0.2253187969326973\n",
      "=====> all mean psnr 27.01232429999652 ssim: 0.9017635122778904 lpips: 0.2627437748014927\n"
     ]
    }
   ],
   "source": [
    "root = '/mnt/new_disk_2/anpei/code/nerf/logs/'\n",
    "root_gt = '/mnt/new_disk2/anpei/code/MVS-NeRF/runs_fine_tuning/'\n",
    "pairs = torch.load('/mnt/new_disk_2/anpei/code/MVS-NeRF/configs/pairs.th')\n",
    "psnr_all,ssim_all,LPIPS_vgg_all = [],[],[]\n",
    "for i_scene, scene in enumerate(['chair','drums','ficus','hotdog','lego','materials','mic','ship']):#\n",
    "    psnr,ssim,LPIPS_vgg = [],[],[]\n",
    "    \n",
    "    files = sorted(glob.glob(f'{root}/{scene}_test/testset_200000/*.png'))\n",
    "    for j, file in enumerate(files):\n",
    "\n",
    "        idx = pairs[f'{scene}_val'][j]\n",
    "        img = cv2.imread(file)[...,::-1]\n",
    "        gt = cv2.imread(f'{root_gt}/{scene}/{scene}/{scene}_00009999_{j:02d}.png')[...,::-1]\n",
    "        gt, img = gt[:,:800].astype('float')/255.0, img.astype('float')/255.0\n",
    "\n",
    "#         H_crop, W_crop = np.array(gt.shape[:2])//10\n",
    "#         img = img[H_crop:-H_crop,W_crop:-W_crop]\n",
    "#         gt = gt[H_crop:-H_crop,W_crop:-W_crop]\n",
    "\n",
    "        psnr.append( mse2psnr(np.mean((gt-img)**2)))\n",
    "        ssim.append( structural_similarity(gt, img, multichannel=True))\n",
    "\n",
    "        img_tensor = torch.from_numpy(img)[None].permute(0,3,1,2).float()*2-1.0 # image should be RGB, IMPORTANT: normalized to [-1,1]\n",
    "        img_gt_tensor = torch.from_numpy(gt)[None].permute(0,3,1,2).float()*2-1.0\n",
    "        LPIPS_vgg.append( loss_fn_vgg(img_tensor, img_gt_tensor).item())\n",
    "\n",
    "    print(f'=====> scene: {scene} mean psnr {np.mean(psnr)} ssim: {np.mean(ssim)} lpips: {np.mean(LPIPS_vgg)}')   \n",
    "    psnr_all.append(psnr);ssim_all.append(ssim);LPIPS_vgg_all.append(LPIPS_vgg)\n",
    "print(f'=====> all mean psnr {np.mean(psnr_all)} ssim: {np.mean(ssim_all)} lpips: {np.mean(LPIPS_vgg_all)}') \n",
    "\n",
    "psnr_all,ssim_all,LPIPS_vgg_all = [],[],[]\n",
    "for i_scene, scene in enumerate(['fern','flower','fortress','horns','leaves', 'orchids', 'room',  'trex']):#,\n",
    "    psnr,ssim,LPIPS_vgg = [],[],[]\n",
    "    \n",
    "    files = sorted(glob.glob(f'{root}/{scene}_test/testset_200000/*.png'))\n",
    "    for j, file in enumerate(files):\n",
    "\n",
    "        idx = pairs[f'{scene}_val'][j]\n",
    "        img = cv2.resize(cv2.imread(file)[...,::-1],(960,640))\n",
    "        gt =  cv2.imread(f'/mnt/new_disk2/anpei/code/MVS-NeRF/results/test3/{scene}_{idx:03d}.png')[...,::-1]\n",
    "        gt, img = gt[:,:960].astype('float')/255.0, img.astype('float')/255.0\n",
    "\n",
    "\n",
    "        psnr.append( mse2psnr(np.mean((gt-img)**2)))\n",
    "        ssim.append( structural_similarity(gt, img, multichannel=True))\n",
    "\n",
    "        img_tensor = torch.from_numpy(img)[None].permute(0,3,1,2).float()*2-1.0 # image should be RGB, IMPORTANT: normalized to [-1,1]\n",
    "        img_gt_tensor = torch.from_numpy(gt)[None].permute(0,3,1,2).float()*2-1.0\n",
    "        LPIPS_vgg.append( loss_fn_vgg(img_tensor, img_gt_tensor).item())\n",
    "\n",
    "    print(f'=====> scene: {scene} mean psnr {np.mean(psnr)} ssim: {np.mean(ssim)} lpips: {np.mean(LPIPS_vgg)}')   \n",
    "    psnr_all.append(psnr);ssim_all.append(ssim);LPIPS_vgg_all.append(LPIPS_vgg)\n",
    "print(f'=====> all mean psnr {np.mean(psnr_all)} ssim: {np.mean(ssim_all)} lpips: {np.mean(LPIPS_vgg_all)}') \n",
    "\n",
    "psnr_all,ssim_all,LPIPS_vgg_all = [],[],[]\n",
    "for i_scene, scene in enumerate([1,8,21,103,114]):#,\n",
    "    psnr,ssim,LPIPS_vgg = [],[],[]\n",
    "    \n",
    "    files = sorted(glob.glob(f'{root}/scan{scene}_test/testset_200000/*.png'))\n",
    "    for j, file in enumerate(files):\n",
    "\n",
    "        idx = pairs[f'dtu_val'][j]\n",
    "        img = cv2.imread(file)[...,::-1]\n",
    "        gt = cv2.imread(f'/mnt/new_disk2/anpei/code/MVS-NeRF/results/test3/scan{scene}_{idx:03d}.png')[...,::-1]\n",
    "        mask = cv2.imread(f'/mnt/new_disk2/anpei/code/MVS-NeRF/results/test3/scan{scene}_{idx:03d}_mask.png')==255\n",
    "        gt, img = gt[:,:640].astype('float')/255.0, img.astype('float')/255.0\n",
    "\n",
    "        gt[mask],img[mask] = 0.0,0.0\n",
    "        psnr.append( mse2psnr(np.mean((gt[~mask]-img[~mask])**2)))\n",
    "        ssim.append( structural_similarity(gt, img, multichannel=True))\n",
    "\n",
    "        img_tensor = torch.from_numpy(img)[None].permute(0,3,1,2).float()*2-1.0 # image should be RGB, IMPORTANT: normalized to [-1,1]\n",
    "        img_gt_tensor = torch.from_numpy(gt)[None].permute(0,3,1,2).float()*2-1.0\n",
    "        LPIPS_vgg.append( loss_fn_vgg(img_tensor, img_gt_tensor).item())\n",
    "\n",
    "    print(f'=====> scene: {scene} mean psnr {np.mean(psnr)} ssim: {np.mean(ssim)} lpips: {np.mean(LPIPS_vgg)}')   \n",
    "    psnr_all.append(psnr);ssim_all.append(ssim);LPIPS_vgg_all.append(LPIPS_vgg)\n",
    "print(f'=====> all mean psnr {np.mean(psnr_all)} ssim: {np.mean(ssim_all)} lpips: {np.mean(LPIPS_vgg_all)}') "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# ours"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "=====> scene: chair mean psnr 26.746873810513947 ssim: 0.9321234340308168 lpips: 0.15475058555603027\n",
      "=====> scene: drums mean psnr 22.28406117543553 ssim: 0.8964614019632969 lpips: 0.21542910858988762\n",
      "=====> scene: ficus mean psnr 26.365789669973488 ssim: 0.9438413276712496 lpips: 0.15914445742964745\n",
      "=====> scene: hotdog mean psnr 32.489636248742805 ssim: 0.9699785599978382 lpips: 0.11295554973185062\n",
      "=====> scene: lego mean psnr 26.79832336361502 ssim: 0.9245065827858229 lpips: 0.18708691000938416\n",
      "=====> scene: materials mean psnr 24.957611270986945 ssim: 0.9249186604651752 lpips: 0.1740873008966446\n",
      "=====> scene: mic mean psnr 29.449610622444368 ssim: 0.9693072200690339 lpips: 0.092950988560915\n",
      "=====> scene: ship mean psnr 26.60832366062154 ssim: 0.8780999869891254 lpips: 0.28621142730116844\n",
      "=====> all mean psnr 26.962528727791707 ssim: 0.9299046467465449 lpips: 0.17282704100944102\n",
      "=====> scene: fern mean psnr 22.61357364768159 ssim: 0.77000724312094 lpips: 0.2827577739953995\n",
      "=====> scene: flower mean psnr 25.52052448547126 ssim: 0.8816617628340061 lpips: 0.19497620686888695\n",
      "=====> scene: fortress mean psnr 28.010539767191638 ssim: 0.8958878940094137 lpips: 0.19440287351608276\n",
      "=====> scene: horns mean psnr 24.99889079936664 ssim: 0.8806854873247154 lpips: 0.2329558990895748\n",
      "=====> scene: leaves mean psnr 21.228759476443667 ssim: 0.8063488185998933 lpips: 0.24219803884625435\n",
      "=====> scene: orchids mean psnr 19.460068266904813 ssim: 0.7087825176683923 lpips: 0.3091204948723316\n",
      "=====> scene: room mean psnr 29.150237317214977 ssim: 0.9570667630128395 lpips: 0.16276488453149796\n",
      "=====> scene: trex mean psnr 24.079128145103503 ssim: 0.8873396265555692 lpips: 0.2028873674571514\n",
      "=====> all mean psnr 24.38271523817226 ssim: 0.8484725141407212 lpips: 0.22775794239714742\n",
      "=====> scene: 1 mean psnr 28.046498782015014 ssim: 0.9341022735743112 lpips: 0.17119702696800232\n",
      "=====> scene: 8 mean psnr 28.875841987497292 ssim: 0.8999499891440654 lpips: 0.26092011481523514\n",
      "=====> scene: 21 mean psnr 24.870061827644722 ssim: 0.9215085855864712 lpips: 0.14212623238563538\n",
      "=====> scene: 103 mean psnr 32.226849694259215 ssim: 0.9638182025610383 lpips: 0.16953438520431519\n",
      "=====> scene: 114 mean psnr 28.46625266665073 ssim: 0.9451965364044997 lpips: 0.15322893112897873\n",
      "=====> all mean psnr 28.497100991613394 ssim: 0.9329151174540771 lpips: 0.17940133810043335\n"
     ]
    }
   ],
   "source": [
    "root = '/mnt/new_disk2/anpei/code/MVS-NeRF/runs_fine_tuning/'\n",
    "pairs = torch.load('/mnt/new_disk_2/anpei/code/MVS-NeRF/configs/pairs.th')\n",
    "psnr_all,ssim_all,LPIPS_vgg_all = [],[],[]\n",
    "for i_scene, scene in enumerate(['chair','drums','ficus','hotdog','lego','materials','mic','ship']):\n",
    "    psnr,ssim,LPIPS_vgg = [],[],[]\n",
    "    \n",
    "    files = sorted(glob.glob(f'{root}/{scene}/{scene}/{scene}_00009999_*'))\n",
    "    for j, file in enumerate(files):\n",
    "\n",
    "        idx = pairs[f'{scene}_val'][j]\n",
    "        img = cv2.imread(file).astype('float')[...,::-1]\n",
    "        gt, img = img[:,:800]/255.0, img[:,800:1600]/255.0\n",
    "\n",
    "#         H_crop, W_crop = np.array(gt.shape[:2])//10\n",
    "#         img = img[H_crop:-H_crop,W_crop:-W_crop]\n",
    "#         gt = gt[H_crop:-H_crop,W_crop:-W_crop]\n",
    "\n",
    "        psnr.append( mse2psnr(np.mean((gt-img)**2)))\n",
    "        ssim.append( structural_similarity(gt, img, multichannel=True))\n",
    "\n",
    "        img_tensor = torch.from_numpy(img)[None].permute(0,3,1,2).float()*2-1.0 # image should be RGB, IMPORTANT: normalized to [-1,1]\n",
    "        img_gt_tensor = torch.from_numpy(gt)[None].permute(0,3,1,2).float()*2-1.0\n",
    "        LPIPS_vgg.append( loss_fn_vgg(img_tensor, img_gt_tensor).item())\n",
    "\n",
    "    print(f'=====> scene: {scene} mean psnr {np.mean(psnr)} ssim: {np.mean(ssim)} lpips: {np.mean(LPIPS_vgg)}')   \n",
    "    psnr_all.append(psnr);ssim_all.append(ssim);LPIPS_vgg_all.append(LPIPS_vgg)\n",
    "print(f'=====> all mean psnr {np.mean(psnr_all)} ssim: {np.mean(ssim_all)} lpips: {np.mean(LPIPS_vgg_all)}') \n",
    "\n",
    "psnr_all,ssim_all,LPIPS_vgg_all = [],[],[]\n",
    "for i_scene, scene in enumerate(['fern','flower','fortress','horns','leaves', 'orchids', 'room',  'trex']):#,\n",
    "    psnr,ssim,LPIPS_vgg = [],[],[]\n",
    "    \n",
    "    files = sorted(glob.glob(f'/mnt/new_disk2/anpei/code/MVS-NeRF/runs_fine_tuning/{scene}/{scene}/{scene}_00009999_*'))\n",
    "    for j, file in enumerate(files):\n",
    "\n",
    "        idx = pairs[f'{scene}_val'][j]\n",
    "        img = cv2.imread(file)[...,::-1]\n",
    "        gt, img = img[:,:960].astype('float')/255.0, img[:,960:960*2].astype('float')/255.0\n",
    "\n",
    "        H_crop, W_crop = np.array(gt.shape[:2])//10\n",
    "        img = img[H_crop:-H_crop,W_crop:-W_crop]\n",
    "        gt = gt[H_crop:-H_crop,W_crop:-W_crop]\n",
    "\n",
    "\n",
    "        psnr.append( mse2psnr(np.mean((gt-img)**2)))\n",
    "        ssim.append( structural_similarity(gt, img, multichannel=True))\n",
    "\n",
    "        img_tensor = torch.from_numpy(img)[None].permute(0,3,1,2).float()*2-1.0 # image should be RGB, IMPORTANT: normalized to [-1,1]\n",
    "        img_gt_tensor = torch.from_numpy(gt)[None].permute(0,3,1,2).float()*2-1.0\n",
    "        LPIPS_vgg.append( loss_fn_vgg(img_tensor, img_gt_tensor).item())\n",
    "\n",
    "    print(f'=====> scene: {scene} mean psnr {np.mean(psnr)} ssim: {np.mean(ssim)} lpips: {np.mean(LPIPS_vgg)}')   \n",
    "    psnr_all.append(psnr);ssim_all.append(ssim);LPIPS_vgg_all.append(LPIPS_vgg)\n",
    "print(f'=====> all mean psnr {np.mean(psnr_all)} ssim: {np.mean(ssim_all)} lpips: {np.mean(LPIPS_vgg_all)}') \n",
    "\n",
    "root = '/mnt/new_disk2/anpei/code/MVS-NeRF/runs_fine_tuning/'\n",
    "psnr_all,ssim_all,LPIPS_vgg_all = [],[],[]\n",
    "for i_scene, scene in enumerate([1,8,21,103,114]):#,\n",
    "    psnr,ssim,LPIPS_vgg = [],[],[]\n",
    "    \n",
    "    files = sorted(glob.glob(f'{root}/dtu_scan{scene}_1h/dtu_scan{scene}_1h/00010239_*'))\n",
    "    for j, file in enumerate(files):\n",
    "\n",
    "        idx = pairs[f'dtu_val'][j]\n",
    "        img = cv2.imread(file)[...,::-1][:,640:1280]\n",
    "        gt = cv2.imread(f'/mnt/new_disk2/anpei/code/MVS-NeRF/results/test3/scan{scene}_{idx:03d}.png')[...,::-1]\n",
    "        mask = cv2.imread(f'/mnt/new_disk2/anpei/code/MVS-NeRF/results/test3/scan{scene}_{idx:03d}_mask.png')==255\n",
    "        gt, img = gt[:,:640].astype('float')/255.0, img.astype('float')/255.0\n",
    "\n",
    "        gt[mask],img[mask] = 0.0,0.0\n",
    "        psnr.append( mse2psnr(np.mean((gt[~mask]-img[~mask])**2)))\n",
    "        ssim.append( structural_similarity(gt, img, multichannel=True))\n",
    "\n",
    "        img_tensor = torch.from_numpy(img)[None].permute(0,3,1,2).float()*2-1.0 # image should be RGB, IMPORTANT: normalized to [-1,1]\n",
    "        img_gt_tensor = torch.from_numpy(gt)[None].permute(0,3,1,2).float()*2-1.0\n",
    "        LPIPS_vgg.append( loss_fn_vgg(img_tensor, img_gt_tensor).item())\n",
    "\n",
    "    print(f'=====> scene: {scene} mean psnr {np.mean(psnr)} ssim: {np.mean(ssim)} lpips: {np.mean(LPIPS_vgg)}')   \n",
    "    psnr_all.append(psnr);ssim_all.append(ssim);LPIPS_vgg_all.append(LPIPS_vgg)\n",
    "print(f'=====> all mean psnr {np.mean(psnr_all)} ssim: {np.mean(ssim_all)} lpips: {np.mean(LPIPS_vgg_all)}') "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# ibrnet"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "=====> scene: fern mean psnr 22.64474646040451 ssim: 0.7736232480476191 lpips: 0.26588304713368416\n",
      "=====> scene: flower mean psnr 26.553349019087786 ssim: 0.9092690161984827 lpips: 0.14575103670358658\n",
      "=====> scene: fortress mean psnr 30.338842953903075 ssim: 0.9368867837660259 lpips: 0.13289865292608738\n",
      "=====> scene: horns mean psnr 25.01290939681414 ssim: 0.9040335882553917 lpips: 0.1899307444691658\n",
      "=====> scene: leaves mean psnr 22.076508076698556 ssim: 0.8430354849586478 lpips: 0.17987846583127975\n",
      "=====> scene: orchids mean psnr 19.007830032899616 ssim: 0.7045611776629173 lpips: 0.2861044891178608\n",
      "=====> scene: room mean psnr 31.05473820815669 ssim: 0.9723299877991765 lpips: 0.08911459799855947\n",
      "=====> scene: trex mean psnr 22.339864946223464 ssim: 0.8421255627008343 lpips: 0.22207806631922722\n",
      "=====> all mean psnr 24.878598636773482 ssim: 0.8607331061736369 lpips: 0.1889548875624314\n",
      "=====> scene: 1 mean psnr 30.99564992655386 ssim: 0.9548394719193786 lpips: 0.1285402663052082\n",
      "=====> scene: 8 mean psnr 32.46173840309124 ssim: 0.9445782574823819 lpips: 0.16979694738984108\n",
      "=====> scene: 21 mean psnr 27.88178725277648 ssim: 0.9467770144187784 lpips: 0.1040295660495758\n",
      "=====> scene: 103 mean psnr 34.399916992709706 ssim: 0.9675776122666799 lpips: 0.15563546121120453\n",
      "=====> scene: 114 mean psnr 31.00119644953211 ssim: 0.9638488318305859 lpips: 0.09874588809907436\n",
      "=====> all mean psnr 31.348057804932676 ssim: 0.955524237583561 lpips: 0.1313496258109808\n"
     ]
    }
   ],
   "source": [
    "# root = '/mnt/new_disk2/anpei/code/IBRNet/logs'\n",
    "pairs = torch.load('/mnt/new_disk_2/anpei/code/MVS-NeRF/configs/pairs.th')\n",
    "# psnr_all,ssim_all,LPIPS_vgg_all = [],[],[]\n",
    "# for i_scene, scene in enumerate(['chair','drums','ficus','hotdog','lego','materials','mic','ship']):#,\n",
    "#     psnr,ssim,LPIPS_vgg = [],[],[]\n",
    "    \n",
    "#     files = sorted(glob.glob(f'{root}/nerf-3view-finetuning-nearest-{scene}/010000_*'))\n",
    "#     for j, file in enumerate(files):\n",
    "\n",
    "#         idx = pairs[f'{scene}_val'][j]\n",
    "#         img = cv2.imread(file).astype('float')[...,::-1]\n",
    "#         gt, img = img[:,800:800*2]/255.0, img[:,800*3:800*4]/255.0\n",
    "\n",
    "# #         H_crop, W_crop = np.array(gt.shape[:2])//10\n",
    "# #         img = img[H_crop:-H_crop,W_crop:-W_crop]\n",
    "# #         gt = gt[H_crop:-H_crop,W_crop:-W_crop]\n",
    "\n",
    "#         psnr.append( mse2psnr(np.mean((gt-img)**2)))\n",
    "#         ssim.append( structural_similarity(gt, img, multichannel=True))\n",
    "\n",
    "#         img_tensor = torch.from_numpy(img)[None].permute(0,3,1,2).float()*2-1.0 # image should be RGB, IMPORTANT: normalized to [-1,1]\n",
    "#         img_gt_tensor = torch.from_numpy(gt)[None].permute(0,3,1,2).float()*2-1.0\n",
    "#         LPIPS_vgg.append( loss_fn_vgg(img_tensor, img_gt_tensor).item())\n",
    "\n",
    "#     print(f'=====> scene: {scene} mean psnr {np.mean(psnr)} ssim: {np.mean(ssim)} lpips: {np.mean(LPIPS_vgg)}')   \n",
    "#     psnr_all.append(psnr);ssim_all.append(ssim);LPIPS_vgg_all.append(LPIPS_vgg)\n",
    "# print(f'=====> all mean psnr {np.mean(psnr_all)} ssim: {np.mean(ssim_all)} lpips: {np.mean(LPIPS_vgg_all)}') \n",
    "\n",
    "psnr_all,ssim_all,LPIPS_vgg_all = [],[],[]\n",
    "for i_scene, scene in enumerate(['fern','flower','fortress','horns','leaves', 'orchids', 'room',  'trex']):#\n",
    "    psnr,ssim,LPIPS_vgg = [],[],[]\n",
    "    \n",
    "    files = sorted(glob.glob(f'/mnt/new_disk2/anpei/code/IBRNet/logs/llff-3view-finetuning-nearest-{scene}/010000_*'))\n",
    "    for j, file in enumerate(files):\n",
    "\n",
    "        idx = pairs[f'{scene}_val'][j]\n",
    "        img = cv2.imread(file)[...,::-1]\n",
    "        gt, img = img[:,1008:1008*2].astype('float')/255.0, img[:,1008*3:1008*4].astype('float')/255.0\n",
    "        img, gt = cv2.resize(img,(960,640)), cv2.resize(gt,(960,640))\n",
    "\n",
    "        H_crop, W_crop = np.array(gt.shape[:2])//10\n",
    "        img = img[H_crop:-H_crop,W_crop:-W_crop]\n",
    "        gt = gt[H_crop:-H_crop,W_crop:-W_crop]\n",
    "\n",
    "        \n",
    "        psnr.append( mse2psnr(np.mean((gt-img)**2)))\n",
    "        ssim.append( structural_similarity(gt, img, multichannel=True))\n",
    "\n",
    "        img_tensor = torch.from_numpy(img)[None].permute(0,3,1,2).float()*2-1.0 # image should be RGB, IMPORTANT: normalized to [-1,1]\n",
    "        img_gt_tensor = torch.from_numpy(gt)[None].permute(0,3,1,2).float()*2-1.0\n",
    "        LPIPS_vgg.append( loss_fn_vgg(img_tensor, img_gt_tensor).item())\n",
    "\n",
    "    print(f'=====> scene: {scene} mean psnr {np.mean(psnr)} ssim: {np.mean(ssim)} lpips: {np.mean(LPIPS_vgg)}')   \n",
    "    psnr_all.append(psnr);ssim_all.append(ssim);LPIPS_vgg_all.append(LPIPS_vgg)\n",
    "print(f'=====> all mean psnr {np.mean(psnr_all)} ssim: {np.mean(ssim_all)} lpips: {np.mean(LPIPS_vgg_all)}') \n",
    "\n",
    "root = '/mnt/new_disk2/anpei/code/IBRNet/logs'\n",
    "psnr_all,ssim_all,LPIPS_vgg_all = [],[],[]\n",
    "for i_scene, scene in enumerate([1,8,21,103,114]):#\n",
    "    psnr,ssim,LPIPS_vgg = [],[],[]\n",
    "    \n",
    "    files = sorted(glob.glob(f'{root}/dtu-3view-finetuning-nearest-scan{scene}/010000_*'))\n",
    "    for j, file in enumerate(files):\n",
    "\n",
    "        idx = pairs[f'dtu_val'][j]\n",
    "        img = cv2.imread(file)[...,::-1][:,3*640:4*640]\n",
    "        gt = cv2.imread(f'/mnt/new_disk2/anpei/code/MVS-NeRF/results/test3/scan{scene}_{idx:03d}.png')[...,::-1]\n",
    "        mask = cv2.imread(f'/mnt/new_disk2/anpei/code/MVS-NeRF/results/test3/scan{scene}_{idx:03d}_mask.png')==255\n",
    "        gt, img = gt[:,:640].astype('float')/255.0, img.astype('float')/255.0\n",
    "\n",
    "        gt[mask],img[mask] = 0.0,0.0\n",
    "        psnr.append( mse2psnr(np.mean((gt[~mask]-img[~mask])**2)))\n",
    "        ssim.append( structural_similarity(gt, img, multichannel=True))\n",
    "\n",
    "        img_tensor = torch.from_numpy(img)[None].permute(0,3,1,2).float()*2-1.0 # image should be RGB, IMPORTANT: normalized to [-1,1]\n",
    "        img_gt_tensor = torch.from_numpy(gt)[None].permute(0,3,1,2).float()*2-1.0\n",
    "        LPIPS_vgg.append( loss_fn_vgg(img_tensor, img_gt_tensor).item())\n",
    "\n",
    "    print(f'=====> scene: {scene} mean psnr {np.mean(psnr)} ssim: {np.mean(ssim)} lpips: {np.mean(LPIPS_vgg)}')   \n",
    "    psnr_all.append(psnr);ssim_all.append(ssim);LPIPS_vgg_all.append(LPIPS_vgg)\n",
    "print(f'=====> all mean psnr {np.mean(psnr_all)} ssim: {np.mean(ssim_all)} lpips: {np.mean(LPIPS_vgg_all)}') "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# pixel nerf"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "=====> scene: chair mean psnr 7.175962813343725 ssim: 0.6243642351905847 lpips: 0.38591109961271286\n",
      "=====> scene: drums mean psnr 8.148548711878252 ssim: 0.6701584468514097 lpips: 0.42121122032403946\n",
      "=====> scene: ficus mean psnr 6.608732738834844 ssim: 0.668716265099144 lpips: 0.3350602239370346\n",
      "=====> scene: hotdog mean psnr 6.799387670799135 ssim: 0.6689815218041557 lpips: 0.43327029794454575\n",
      "=====> scene: lego mean psnr 7.740217521658803 ssim: 0.6710903029993184 lpips: 0.42670799791812897\n",
      "=====> scene: materials mean psnr 7.609290420358684 ssim: 0.6441046576733512 lpips: 0.43245941400527954\n",
      "=====> scene: mic mean psnr 7.707203698223274 ssim: 0.7294597852809476 lpips: 0.32929887622594833\n",
      "=====> scene: ship mean psnr 7.295484760785579 ssim: 0.5836685948507447 lpips: 0.5257005095481873\n",
      "=====> all mean psnr 7.385603541985287 ssim: 0.657567976218707 lpips: 0.4112024549394846\n",
      "=====> scene: fern mean psnr 12.397648684821284 ssim: 0.5312397318110376 lpips: 0.6500117480754852\n",
      "=====> scene: flower mean psnr 9.99675489427281 ssim: 0.43323656344453193 lpips: 0.7075561136007309\n",
      "=====> scene: fortress mean psnr 14.073488262986546 ssim: 0.6736649368929569 lpips: 0.6075489073991776\n",
      "=====> scene: horns mean psnr 11.71002161685521 ssim: 0.5157559834106309 lpips: 0.705190509557724\n",
      "=====> scene: leaves mean psnr 9.847068575637605 ssim: 0.2681360856716045 lpips: 0.6947661340236664\n",
      "=====> scene: orchids mean psnr 9.624184850140201 ssim: 0.3168018322171233 lpips: 0.7207814902067184\n",
      "=====> scene: room mean psnr 11.750716240417157 ssim: 0.6906632306577324 lpips: 0.6112178564071655\n",
      "=====> scene: trex mean psnr 10.55211637013632 ssim: 0.45770187915676297 lpips: 0.6672322899103165\n",
      "=====> all mean psnr 11.243999936908391 ssim: 0.48590003040779756 lpips: 0.6705381311476231\n",
      "=====> scene: 1 mean psnr 21.64330896303546 ssim: 0.8268906240371169 lpips: 0.3728666678071022\n",
      "=====> scene: 8 mean psnr 23.69860813191727 ssim: 0.8291247579664462 lpips: 0.38369619846343994\n",
      "=====> scene: 21 mean psnr 16.03916045790063 ssim: 0.6905614284959526 lpips: 0.4074021577835083\n",
      "=====> scene: 103 mean psnr 16.75554527345504 ssim: 0.8356182752116726 lpips: 0.3762983977794647\n",
      "=====> scene: 114 mean psnr 18.403335481905653 ssim: 0.7632868039029544 lpips: 0.371926873922348\n",
      "=====> all mean psnr 19.30799166164281 ssim: 0.7890963779228286 lpips: 0.3824380591511726\n"
     ]
    }
   ],
   "source": [
    "root = '/mnt/new_disk2/anpei/code/pixel-nerf/visuals/dtu'\n",
    "root_gt = '/mnt/new_disk2/anpei/code/MVS-NeRF/runs_fine_tuning/'\n",
    "pairs = torch.load('/mnt/new_disk_2/anpei/code/MVS-NeRF/configs/pairs.th')\n",
    "psnr_all,ssim_all,LPIPS_vgg_all = [],[],[]\n",
    "for i_scene, scene in enumerate(['chair','drums','ficus','hotdog','lego','materials','mic','ship']):#,\n",
    "    psnr,ssim,LPIPS_vgg = [],[],[]\n",
    "    \n",
    "    pairs_idx = pairs[f'{scene}_val']\n",
    "    for j, file in enumerate(pairs_idx):\n",
    "\n",
    "        idx = pairs_idx[j]\n",
    "        img = cv2.imread(f'{root}/{scene}_{idx:03d}.png')[...,::-1]\n",
    "        gt = cv2.imread(f'{root_gt}/{scene}/{scene}/{scene}_00009999_{j:02d}.png')[...,::-1]\n",
    "        gt, img = gt[:,:800].astype('float')/255.0, img.astype('float')/255.0\n",
    "\n",
    "#         H_crop, W_crop = np.array(gt.shape[:2])//10\n",
    "#         img = img[H_crop:-H_crop,W_crop:-W_crop]\n",
    "#         gt = gt[H_crop:-H_crop,W_crop:-W_crop]\n",
    "\n",
    "        psnr.append( mse2psnr(np.mean((gt-img)**2)))\n",
    "        ssim.append( structural_similarity(gt, img, multichannel=True))\n",
    "\n",
    "        img_tensor = torch.from_numpy(img)[None].permute(0,3,1,2).float()*2-1.0 # image should be RGB, IMPORTANT: normalized to [-1,1]\n",
    "        img_gt_tensor = torch.from_numpy(gt)[None].permute(0,3,1,2).float()*2-1.0\n",
    "        LPIPS_vgg.append( loss_fn_vgg(img_tensor, img_gt_tensor).item())\n",
    "\n",
    "    print(f'=====> scene: {scene} mean psnr {np.mean(psnr)} ssim: {np.mean(ssim)} lpips: {np.mean(LPIPS_vgg)}')   \n",
    "    psnr_all.append(psnr);ssim_all.append(ssim);LPIPS_vgg_all.append(LPIPS_vgg)\n",
    "print(f'=====> all mean psnr {np.mean(psnr_all)} ssim: {np.mean(ssim_all)} lpips: {np.mean(LPIPS_vgg_all)}') \n",
    "\n",
    "psnr_all,ssim_all,LPIPS_vgg_all = [],[],[]\n",
    "for i_scene, scene in enumerate(['fern','flower','fortress','horns','leaves', 'orchids', 'room',  'trex']):#,\n",
    "    psnr,ssim,LPIPS_vgg = [],[],[]\n",
    "    \n",
    "    pairs_idx = pairs[f'{scene}_val']\n",
    "    for j, file in enumerate(pairs_idx):\n",
    "\n",
    "        idx = pairs_idx[j]\n",
    "        img = cv2.resize(cv2.imread(f'{root}/{scene}_{idx:03d}.png')[...,::-1],(960,640))\n",
    "        gt =  cv2.imread(f'/mnt/new_disk2/anpei/code/MVS-NeRF/results/test3/{scene}_{idx:03d}.png')[...,::-1]\n",
    "        gt, img = gt[:,:960].astype('float')/255.0, img.astype('float')/255.0\n",
    "\n",
    "\n",
    "        psnr.append( mse2psnr(np.mean((gt-img)**2)))\n",
    "        ssim.append( structural_similarity(gt, img, multichannel=True))\n",
    "\n",
    "        img_tensor = torch.from_numpy(img)[None].permute(0,3,1,2).float()*2-1.0 # image should be RGB, IMPORTANT: normalized to [-1,1]\n",
    "        img_gt_tensor = torch.from_numpy(gt)[None].permute(0,3,1,2).float()*2-1.0\n",
    "        LPIPS_vgg.append( loss_fn_vgg(img_tensor, img_gt_tensor).item())\n",
    "\n",
    "    print(f'=====> scene: {scene} mean psnr {np.mean(psnr)} ssim: {np.mean(ssim)} lpips: {np.mean(LPIPS_vgg)}')   \n",
    "    psnr_all.append(psnr);ssim_all.append(ssim);LPIPS_vgg_all.append(LPIPS_vgg)\n",
    "print(f'=====> all mean psnr {np.mean(psnr_all)} ssim: {np.mean(ssim_all)} lpips: {np.mean(LPIPS_vgg_all)}') \n",
    "\n",
    "psnr_all,ssim_all,LPIPS_vgg_all,depth_acc = [],[],[],[]\n",
    "for i_scene, scene in enumerate([1,8,21,103,114]):#,\n",
    "    psnr,ssim,LPIPS_vgg = [],[],[]\n",
    "    \n",
    "    pairs_idx = pairs[f'dtu_val']\n",
    "    for j, file in enumerate(pairs_idx):\n",
    "\n",
    "        idx = pairs_idx[j]\n",
    "        img = cv2.imread(f'{root}/scan{scene}_{idx:03d}.png')[...,::-1]\n",
    "        \n",
    "        gt = cv2.imread(f'/mnt/new_disk2/anpei/code/MVS-NeRF/results/test3/scan{scene}_{idx:03d}.png')[...,::-1]\n",
    "        mask = cv2.imread(f'/mnt/new_disk2/anpei/code/MVS-NeRF/results/test3/scan{scene}_{idx:03d}_mask.png')==255\n",
    "        gt, img = gt[:,:640].astype('float')/255.0, img.astype('float')/255.0\n",
    "\n",
    "        gt[mask],img[mask] = 0.0,0.0\n",
    "        psnr.append( mse2psnr(np.mean((gt[~mask]-img[~mask])**2)))\n",
    "        ssim.append( structural_similarity(gt, img, multichannel=True))\n",
    "\n",
    "        img_tensor = torch.from_numpy(img)[None].permute(0,3,1,2).float()*2-1.0 # image should be RGB, IMPORTANT: normalized to [-1,1]\n",
    "        img_gt_tensor = torch.from_numpy(gt)[None].permute(0,3,1,2).float()*2-1.0\n",
    "        LPIPS_vgg.append( loss_fn_vgg(img_tensor, img_gt_tensor).item())\n",
    "        \n",
    "        # depth\n",
    "#         depth_pred = torch.load(f'{root}/scan{scene}_{idx:03d}_depth.th')\n",
    "#         depth_gt,_ =  read_depth(f'/mnt/data/new_disk/sungx/data/mvs_dataset/DTU/mvs_training/dtu/Depths/scan{scene}/depth_map_{idx:04d}.pfm')\n",
    "        \n",
    "#         mask_gt = depth_gt>0\n",
    "#         abs_err = abs_error(depth_pred*1.5, depth_gt/200, mask_gt).numpy()\n",
    "\n",
    "#         eval_metric = [0.01,0.05, 0.1]\n",
    "#         depth_acc[f'abs_err'][f'{scene}'] = np.mean(abs_err)\n",
    "#         depth_acc[f'acc_l_{eval_metric[0]}'][f'{scene}'] = acc_threshold(abs_err,eval_metric[0]).mean()\n",
    "#         depth_acc[f'acc_l_{eval_metric[1]}'][f'{scene}'] = acc_threshold(abs_err,eval_metric[1]).mean()\n",
    "#         depth_acc[f'acc_l_{eval_metric[2]}'][f'{scene}'] = acc_threshold(abs_err,eval_metric[2]).mean()\n",
    "\n",
    "    print(f'=====> scene: {scene} mean psnr {np.mean(psnr)} ssim: {np.mean(ssim)} lpips: {np.mean(LPIPS_vgg)}')   \n",
    "    psnr_all.append(psnr);ssim_all.append(ssim);LPIPS_vgg_all.append(LPIPS_vgg)\n",
    "print(f'=====> all mean psnr {np.mean(psnr_all)} ssim: {np.mean(ssim_all)} lpips: {np.mean(LPIPS_vgg_all)}') "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}