renderer_video.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# import packages"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys,os,imageio\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",
    "from scipy.spatial.transform import Rotation as R\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(0)\n",
    "os.environ['CUDA_VISIBLE_DEVICES'] = '0'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Rendering video from finetuned ckpts"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "metadata": {},
   "outputs": [],
   "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",
    "\n",
    "def normalize(x):\n",
    "    return x / np.linalg.norm(x, axis=-1, keepdims=True)\n",
    "\n",
    "def viewmatrix(z, up, pos):\n",
    "    vec2 = normalize(z)\n",
    "    vec1_avg = up\n",
    "    vec0 = normalize(np.cross(vec1_avg, vec2))\n",
    "    vec1 = normalize(np.cross(vec2, vec0))\n",
    "    m = np.eye(4)\n",
    "    m[:3] = np.stack([vec0, vec1, vec2, pos], 1)\n",
    "    return m\n",
    "\n",
    "def ptstocam(pts, c2w):\n",
    "    tt = np.matmul(c2w[:3,:3].T, (pts-c2w[:3,3])[...,np.newaxis])[...,0]\n",
    "    return tt\n",
    "\n",
    "def poses_avg(poses):\n",
    "\n",
    "    center = poses[:, :3, 3].mean(0)\n",
    "    print('center in poses_avg', center)\n",
    "    vec2 = normalize(poses[:, :3, 2].sum(0))\n",
    "    up = poses[:, :3, 1].sum(0)\n",
    "    c2w = viewmatrix(vec2, up, center)\n",
    "    \n",
    "    return c2w\n",
    "\n",
    "def render_path_imgs(c2ws_all, focal):\n",
    "    T = c2ws_all[...,3]\n",
    "\n",
    "    return render_poses\n",
    "\n",
    "def render_path_spiral(c2w, up, rads, focal, zdelta, zrate, N_rots=2, N=120):\n",
    "    render_poses = []\n",
    "    rads = np.array(list(rads) + [1.])\n",
    "    \n",
    "    for theta in np.linspace(0., 2. * np.pi * N_rots, N+1)[:-1]:\n",
    "    # for theta in np.linspace(0., 5, N+1)[:-1]:\n",
    "        # spiral\n",
    "        # c = np.dot(c2w[:3,:4], np.array([np.cos(theta), -np.sin(theta), -np.sin(theta*zrate), 1.]) * rads)\n",
    "\n",
    "        # 关于使用平均姿态的相关问题\n",
    "        # 从训练集推算出来的平均姿态方向基本平行于z轴，因为训练集中大多数图片是正面的，\n",
    "        # 但是存在一个问题，将z和平均姿态相乘后得到的方向也基本上和z平行，所以无论怎么调整看起来都是平行的，\n",
    "        # 别用平均姿态看其他位置的照片，直接用世界坐标系即可！！！！\n",
    "        # 但是需要用别的姿态大致估计一下位置参数\n",
    "        c = np.array([(np.cos(theta)*theta)/10, (-np.sin(theta)*theta)/10, -0.1]) \n",
    "\n",
    "        # 这个是因为作者在读取并规范化相机姿态的时候作了poses*blender2opencv，转换了坐标系，\n",
    "        # 我用的数据无需转换，但是这里加个负号就解决了，目前不影响什么，记住就行\n",
    "        z = -(normalize(c - np.array([0,0,-focal])))\n",
    "        print(\"c\", c)\n",
    "        print(\"z\", z)\n",
    "        render_poses.append(viewmatrix(z, up, c))\n",
    "    return render_poses\n",
    "\n",
    "def get_spiral(c2ws_all, near_far, rads_scale=0.5, N_views=120):\n",
    "\n",
    "    # center pose\n",
    "    c2w = poses_avg(c2ws_all)\n",
    "    print('poses_avg', c2w)\n",
    "    \n",
    "    # Get average pose\n",
    "    up = normalize(c2ws_all[:, :3, 1].sum(0))\n",
    "\n",
    "    # Find a reasonable \"focus depth\" for this dataset\n",
    "    close_depth, inf_depth = near_far\n",
    "    print('near and far bounds', close_depth, inf_depth)\n",
    "    dt = .75\n",
    "    mean_dz = 1./(((1.-dt)/close_depth + dt/inf_depth))\n",
    "    focal = mean_dz\n",
    "    print(focal)\n",
    "\n",
    "    # Get radii for spiral path\n",
    "    shrink_factor = .8\n",
    "    zdelta = close_depth * .2\n",
    "    tt = c2ws_all[:,:3,3] - c2w[:3,3][None]\n",
    "    rads = np.percentile(np.abs(tt), 70, 0)*rads_scale\n",
    "    print(\"rads\",rads)\n",
    "    render_poses = render_path_spiral(c2w, up, rads, focal, zdelta, zrate=.5, N=N_views)\n",
    "    return np.stack(render_poses)\n",
    "\n",
    "def position2angle(position, N_views=16, N_rots = 2):\n",
    "    ''' nx3 '''\n",
    "    position = normalize(position)\n",
    "    theta = np.arccos(position[:,2])/np.pi*180\n",
    "    phi = np.arctan2(position[:,1],position[:,0])/np.pi*180\n",
    "    return [theta,phi]\n",
    "\n",
    "def pose_spherical_nerf(euler, radius=0.01):\n",
    "    c2ws_render = np.eye(4)\n",
    "    c2ws_render[:3,:3] =  R.from_euler('xyz', euler, degrees=True).as_matrix()\n",
    "    # 保留旋转矩阵的最后一列再乘个系数就能当作位置？\n",
    "    c2ws_render[:3,3]  = c2ws_render[:3,:3] @ np.array([0.0,0.0,-radius])\n",
    "    return c2ws_render\n",
    "\n",
    "def create_spheric_poses(radius, n_poses=120):\n",
    "    \"\"\"\n",
    "    Create circular poses around z axis.\n",
    "    Inputs:\n",
    "        radius: the (negative) height and the radius of the circle.\n",
    "    Outputs:\n",
    "        spheric_poses: (n_poses, 3, 4) the poses in the circular path\n",
    "    \"\"\"\n",
    "\n",
    "    def spheric_pose(theta, phi, radius):\n",
    "        trans_t = lambda t: np.array([\n",
    "            [1, 0, 0, 0],\n",
    "            [0, 1, 0, -0.9 * t],\n",
    "            [0, 0, 1, t],\n",
    "            [0, 0, 0, 1],\n",
    "        ])\n",
    "\n",
    "        rot_phi = lambda phi: np.array([\n",
    "            [1, 0, 0, 0],\n",
    "            [0, np.cos(phi), -np.sin(phi), 0],\n",
    "            [0, np.sin(phi), np.cos(phi), 0],\n",
    "            [0, 0, 0, 1],\n",
    "        ])\n",
    "\n",
    "        rot_theta = lambda th: np.array([\n",
    "            [np.cos(th), 0, -np.sin(th), 0],\n",
    "            [0, 1, 0, 0],\n",
    "            [np.sin(th), 0, np.cos(th), 0],\n",
    "            [0, 0, 0, 1],\n",
    "        ])\n",
    "\n",
    "        c2w = rot_theta(theta) @ rot_phi(phi) @ trans_t(radius)\n",
    "        c2w = np.array([[-1, 0, 0, 0], [0, 0, 1, 0], [0, 1, 0, 0], [0, 0, 0, 1]]) @ c2w\n",
    "        return c2w[:3]\n",
    "\n",
    "    spheric_poses = []\n",
    "    for th in np.linspace(0, 2 * np.pi, n_poses + 1)[:-1]:\n",
    "        spheric_poses += [spheric_pose(th, -np.pi / 5, radius)]  # 36 degree view downwards\n",
    "    return np.stack(spheric_poses, 0)\n",
    "\n",
    "def nerf_video_path(c2ws, theta_range=10,phi_range=20,N_views=120):\n",
    "    c2ws = torch.tensor(c2ws)\n",
    "    mean_position = torch.mean(c2ws[:,:3, 3],dim=0).reshape(1,3).cpu().numpy()\n",
    "    rotvec = []\n",
    "    for i in range(c2ws.shape[0]):\n",
    "        r = R.from_matrix(c2ws[i, :3, :3])\n",
    "        euler_ange = r.as_euler('xyz', degrees=True).reshape(1, 3)\n",
    "        if i:\n",
    "            mask = np.abs(euler_ange - rotvec[0])>180\n",
    "            euler_ange[mask] += 360.0\n",
    "        rotvec.append(euler_ange)\n",
    "    # 采用欧拉角做平均的方法求旋转矩阵的平均\n",
    "    rotvec = np.mean(np.stack(rotvec), axis=0)\n",
    "#     render_poses = [pose_spherical_nerf(rotvec)]\n",
    "    render_poses = [pose_spherical_nerf(rotvec+np.array([angle,0.0,-phi_range])) for angle in np.linspace(-theta_range,theta_range,N_views//4, endpoint=False)]\n",
    "    render_poses += [pose_spherical_nerf(rotvec+np.array([theta_range,0.0,angle])) for angle in np.linspace(-phi_range,phi_range,N_views//4, endpoint=False)]\n",
    "    render_poses += [pose_spherical_nerf(rotvec+np.array([angle,0.0,phi_range])) for angle in np.linspace(theta_range,-theta_range,N_views//4, endpoint=False)]\n",
    "    render_poses += [pose_spherical_nerf(rotvec+np.array([-theta_range,0.0,angle])) for angle in np.linspace(phi_range,-phi_range,N_views//4, endpoint=False)]\n",
    "    # render_poses = torch.from_numpy(np.stack(render_poses)).float().to(device)\n",
    "    return render_poses\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### LLFF video rendering"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Found ckpts ['./runs_fine_tuning/xgaze_11images_cropped_colmapCODE-ft/ckpts/latest.tar']\n",
      "Reloading from ./runs_fine_tuning/xgaze_11images_cropped_colmapCODE-ft/ckpts/latest.tar\n",
      "11 11 /home/hengfei/Desktop/research/mvsnerf/xgaze/xgaze_11images_cropped_colmapCODE\n",
      "center in center_poses [ 2.39707244e-17 -1.21115239e-16 -2.01858732e-17]\n",
      "scale_factor 3.8620493751040295\n",
      "===> valing index: [0 5]\n",
      "original focal 13230.734319550395\n",
      "porcessed focal [3175.3762366920946, 2116.9174911280634]\n",
      "center in center_poses [ 2.39707244e-17 -1.21115239e-16 -2.01858732e-17]\n",
      "====> ref idx: [1 7 6]\n",
      "center in poses_avg [-0.14081041  0.08933562 -0.00931392]\n",
      "poses_avg [[ 9.96576197e-01  4.14913612e-04  8.26783576e-02 -1.40810414e-01]\n",
      " [ 5.04793126e-03 -9.98427221e-01 -5.58354958e-02  8.93356208e-02]\n",
      " [ 8.25251559e-02  5.60616808e-02 -9.95010898e-01 -9.31391948e-03]\n",
      " [ 0.00000000e+00  0.00000000e+00  0.00000000e+00  1.00000000e+00]]\n",
      "near and far bounds 1.3106866440139757 2.1600577335631135\n",
      "1.8588996320044613\n",
      "rads [9.12077768 7.36954194 1.10293058]\n",
      "c [ 0.  -0.  -0.1]\n",
      "z [-0.  0. -1.]\n",
      "c [ 0.0382665  -0.01703734 -0.1       ]\n",
      "z [-0.02174977  0.00968362 -0.99971655]\n",
      "c [ 0.05605695 -0.06225756 -0.1       ]\n",
      "z [-0.03183438  0.03535566 -0.99886763]\n",
      "c [ 0.03883222 -0.11951329 -0.1       ]\n",
      "z [-0.02202143  0.06777501 -0.99745757]\n",
      "c [-0.01751391 -0.16663374 -0.1       ]\n",
      "z [ 0.00991244  0.09431056 -0.99549348]\n",
      "c [-0.10471976 -0.18137994 -0.1       ]\n",
      "z [ 0.05911944  0.10239788 -0.99298518]\n",
      "c [-0.20332815 -0.14772655 -0.1       ]\n",
      "z [ 0.11443729  0.08314355 -0.98994508]\n",
      "c [-0.28680786 -0.06096289 -0.1       ]\n",
      "z [ 0.16084137  0.03418789 -0.98638798]\n",
      "c [-0.32778041  0.06967188 -0.1       ]\n",
      "z [ 0.18306265 -0.03891117 -0.9823309 ]\n",
      "c [-0.30499222  0.22158982 -0.1       ]\n",
      "z [ 0.16954874 -0.12318437 -0.97779284]\n",
      "c [-0.20943951  0.36275987 -0.1       ]\n",
      "z [ 0.1158347  -0.20063159 -0.97279457]\n",
      "c [-0.04816326  0.45824279 -0.1       ]\n",
      "z [ 0.02648879 -0.25202405 -0.96735837]\n",
      "c [ 0.15532888  0.47805315 -0.1       ]\n",
      "z [-0.084911   -0.26132919 -0.96150776]\n",
      "c [ 0.3643702   0.40467411 -0.1       ]\n",
      "z [-0.1978913  -0.21978055 -0.95526726]\n",
      "c [ 0.53573104  0.23852283 -0.1       ]\n",
      "z [-0.28894642 -0.12864723 -0.94866214]\n",
      "c [ 6.28318531e-01  1.53893655e-16 -1.00000000e-01]\n",
      "z [-3.36402916e-01 -8.23949509e-17 -9.41718152e-01]\n",
      "c [ 0.61226404 -0.27259751 -0.1       ]\n",
      "z [-0.32528125  0.14482454 -0.93446132]\n",
      "c [ 0.47648411 -0.52918922 -0.1       ]\n",
      "z [-0.25110106  0.27887598 -0.92691771]\n",
      "c [ 0.23299332 -0.71707972 -0.1       ]\n",
      "z [-0.12175069  0.3747101  -0.91911322]\n",
      "c [-0.08319108 -0.79151028 -0.1       ]\n",
      "z [ 0.04309125  0.40998585 -0.9110734 ]\n",
      "c [-0.41887902 -0.72551975 -0.1       ]\n",
      "z [ 0.21500587  0.37240108 -0.9028233 ]\n",
      "c [-0.71164852 -0.51704291 -0.1       ]\n",
      "z [ 0.36186794  0.26291245 -0.8943873 ]\n",
      "c [-0.90139612 -0.19159766 -0.1       ]\n",
      "z [ 0.45394675  0.09648936 -0.885789  ]\n",
      "c [-0.94236867  0.20030664 -0.1       ]\n",
      "z [ 0.46989919 -0.09988016 -0.87705114]\n",
      "c [-0.81331259  0.59090619 -0.1       ]\n",
      "z [ 0.40145231 -0.29167217 -0.86819548]\n",
      "c [-0.52359878  0.90689968 -0.1       ]\n",
      "z [ 0.25578404 -0.44303095 -0.85924275]\n",
      "c [-0.11384043  1.08311933 -0.1       ]\n",
      "z [ 0.05502791 -0.52355558 -0.85021261]\n",
      "c [ 0.34948999  1.07561958 -0.1       ]\n",
      "z [-0.16712965 -0.51437218 -0.84112362]\n",
      "c [ 0.78479736  0.87160577 -0.1       ]\n",
      "z [-0.37122418 -0.41228622 -0.8319932 ]\n",
      "c [ 1.10972858  0.494083   -0.1       ]\n",
      "z [-0.51914643 -0.23113888 -0.82283765]\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  0%|          | 0/30 [00:02<?, ?it/s]\n"
     ]
    },
    {
     "ename": "KeyboardInterrupt",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
      "\u001b[0;32m/tmp/ipykernel_1747469/2421088467.py\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m    106\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    107\u001b[0m                 \u001b[0;31m# rendering\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 108\u001b[0;31m                 rgb, disp, acc, depth_pred, alpha, extras = rendering(args, pose_source, xyz_coarse_sampled,\n\u001b[0m\u001b[1;32m    109\u001b[0m                                                                        \u001b[0mxyz_NDC\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mz_vals\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mrays_o\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mrays_d\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    110\u001b[0m                                                                        volume_feature,imgs_source, **render_kwargs_train)\n",
      "\u001b[0;32m~/Desktop/research/mvsnerf/renderer.py\u001b[0m in \u001b[0;36mrendering\u001b[0;34m(args, pose_ref, rays_pts, rays_ndc, depth_candidates, rays_o, rays_dir, volume_feature, imgs, network_fn, img_feat, network_query_fn, white_bkgd, **kwargs)\u001b[0m\n\u001b[1;32m    158\u001b[0m         \u001b[0minput_feat\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcat\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0minput_feat\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m...\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;36m8\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0mraw\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m...\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;36m4\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdim\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;34m-\u001b[0m\u001b[0;36m1\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    159\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 160\u001b[0;31m     \u001b[0mdists\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mdepth2dist\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mdepth_candidates\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mcos_angle\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    161\u001b[0m     \u001b[0;31m# dists = ndc2dist(rays_ndc)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    162\u001b[0m     \u001b[0mrgb_map\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdisp_map\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0macc_map\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mweights\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdepth_map\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0malpha\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mraw2outputs\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mraw\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdepth_candidates\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdists\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mwhite_bkgd\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mnet_type\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/Desktop/research/mvsnerf/renderer.py\u001b[0m in \u001b[0;36mdepth2dist\u001b[0;34m(z_vals, cos_angle)\u001b[0m\n\u001b[1;32m      7\u001b[0m     \u001b[0mdevice\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mz_vals\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdevice\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      8\u001b[0m     \u001b[0mdists\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mz_vals\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m...\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;36m1\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m]\u001b[0m \u001b[0;34m-\u001b[0m \u001b[0mz_vals\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m...\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m:\u001b[0m\u001b[0;34m-\u001b[0m\u001b[0;36m1\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 9\u001b[0;31m     \u001b[0mdists\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcat\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mdists\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mTensor\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;36m1e10\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mto\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mdevice\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mexpand\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mdists\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m...\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m:\u001b[0m\u001b[0;36m1\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mshape\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m-\u001b[0m\u001b[0;36m1\u001b[0m\u001b[0;34m)\u001b[0m  \u001b[0;31m# [N_rays, N_samples]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     10\u001b[0m     \u001b[0mdists\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mdists\u001b[0m \u001b[0;34m*\u001b[0m \u001b[0mcos_angle\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0munsqueeze\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m-\u001b[0m\u001b[0;36m1\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     11\u001b[0m     \u001b[0;32mreturn\u001b[0m \u001b[0mdists\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m: "
     ]
    }
   ],
   "source": [
    "\n",
    "for i_scene, scene in enumerate(['xgaze_11images_cropped_colmapCODE']):#'horns','flower','orchids', 'room','leaves','fern','trex','fortress'\n",
    "    # add --use_color_volume if the ckpts are fintuned with this flag\n",
    "    cmd = f'--datadir /home/hengfei/Desktop/research/mvsnerf/xgaze/{scene}  \\\n",
    "     --dataset_name llff --imgScale_test {1.0}  --netwidth 128 --net_type v0 '# Please check whether finetuning setting is same with rendering setting, especially on use_color_volume, pad, and use_disp\n",
    "\n",
    "    is_finetued = True # set False if rendering without finetuning\n",
    "    if is_finetued:\n",
    "        cmd += f'--ckpt ./runs_fine_tuning/{scene}-ft/ckpts/latest.tar'\n",
    "    else:\n",
    "        cmd += '--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",
    "#     args.use_color_volume = False if not is_finetued else args.use_color_volume\n",
    "\n",
    "    # create models\n",
    "    if i_scene==0 or is_finetued:\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 #the padding value should be same as your finetuning ckpt\n",
    "    args.chunk = 5120\n",
    "\n",
    "\n",
    "    dataset = dataset_dict[args.dataset_name](args, split=datatype)\n",
    "    val_idx = dataset.img_idx\n",
    "    \n",
    "\n",
    "    save_dir = f'/home/hengfei/Desktop/research/mvsnerf/results/videos'\n",
    "    os.makedirs(save_dir, exist_ok=True)\n",
    "    MVSNet.train()\n",
    "    MVSNet = MVSNet.cuda()\n",
    "    \n",
    "    with torch.no_grad():\n",
    "\n",
    "        c2ws_all = dataset.poses\n",
    "\n",
    "        if is_finetued:   \n",
    "            # large baselien\n",
    "            imgs_source, proj_mats, near_far_source, pose_source = dataset.read_source_views(device=device)\n",
    "            volume_feature = torch.load(args.ckpt)['volume']['feat_volume']\n",
    "            volume_feature = RefVolume(volume_feature.detach()).cuda()\n",
    "            \n",
    "            pad *= args.imgScale_test\n",
    "            w2cs, c2ws = pose_source['w2cs'], pose_source['c2ws']\n",
    "            pair_idx = torch.load('configs/pairs.th')[f'{scene}_train']\n",
    "            # 为啥这个地方只用训练集的相机姿态？\n",
    "            c2ws_render = get_spiral(c2ws_all[pair_idx], near_far_source, rads_scale = 10, N_views=30)# you can enlarge the rads_scale if you want to render larger baseline\n",
    "            # c2ws_render = nerf_video_path(c2ws_all[pair_idx], N_views=40)# you can enlarge the rads_scale if you want to render larger baseline\n",
    "            # c2ws_render = create_spheric_poses(2, n_poses=10)\n",
    "            # c2ws_render = render_path_imgs(c2ws_all[pair_idx])\n",
    "        else:            \n",
    "            # neighboring views with position distance\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, lindisp=args.use_disp)\n",
    "            \n",
    "            pad *= args.imgScale_test\n",
    "            w2cs, c2ws = pose_source['w2cs'], pose_source['c2ws']\n",
    "            pair_idx = torch.load('configs/pairs.th')[f'{scene}_train']\n",
    "            c2ws_render = get_spiral(c2ws_all[pair_idx], near_far_source, rads_scale = 0.1, N_views=60)# you can enlarge the rads_scale if you want to render larger baseline\n",
    "            \n",
    "        c2ws_render = torch.from_numpy(np.stack(c2ws_render)).float().to(device)\n",
    "\n",
    "            \n",
    "        imgs_source = unpreprocess(imgs_source)\n",
    "\n",
    "        try:\n",
    "            tqdm._instances.clear() \n",
    "        except Exception:     \n",
    "            pass\n",
    "        \n",
    "        frames = []\n",
    "        img_directions = dataset.directions.to(device)\n",
    "        for i, c2w in enumerate(tqdm(c2ws_render)):\n",
    "            torch.cuda.empty_cache()\n",
    "            \n",
    "            rays_o, rays_d = get_rays(img_directions, c2w)  # both (h*w, 3)\n",
    "            rays = torch.cat([rays_o, rays_d,\n",
    "                     near_far_source[0] * torch.ones_like(rays_o[:, :1]),\n",
    "                     near_far_source[1] * torch.ones_like(rays_o[:, :1])],\n",
    "                    1).to(device)  # (H*W, 3)\n",
    "            \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=args.use_disp)\n",
    "\n",
    "                # Converting world coordinate to ndc coordinate\n",
    "                H, W = imgs_source.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, lindisp=args.use_disp)\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",
    "                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",
    "            H_crop, W_crop = np.array(rgb_rays.shape[:2])//20\n",
    "            rgb_rays = rgb_rays[H_crop:-H_crop,W_crop:-W_crop]\n",
    "            depth_rays_preds = depth_rays_preds[H_crop:-H_crop,W_crop:-W_crop]\n",
    "            # print(depth_rays_preds)\n",
    "            # img_vis = np.concatenate((rgb_rays*255,depth_rays_preds),axis=1)\n",
    "            img_vis = rgb_rays*255\n",
    "            \n",
    "            imageio.imwrite(f'{save_dir}/video_imgs_1/{str(i).zfill(4)}.png', img_vis.astype('uint8'))\n",
    "            frames.append(img_vis.astype('uint8'))\n",
    "\n",
    "\n",
    "    imageio.mimwrite(f'{save_dir}/ft_{scene}_spiral_test_v4.mp4', np.stack(frames), fps=10, quality=10)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### NeRF video rendering"
   ]
  },
  {
   "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",
      "===> valing index: [26 60 13 47]\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/anpei/anaconda3/lib/python3.7/site-packages/ipykernel_launcher.py:57: FutureWarning: The input object of type 'Tensor' is an array-like implementing one of the corresponding protocols (`__array__`, `__array_interface__` or `__array_struct__`); but not a sequence (or 0-D). In the future, this object will be coerced as if it was first converted using `np.array(obj)`. To retain the old behaviour, you have to either modify the type 'Tensor', or assign to an empty array created with `np.empty(correct_shape, dtype=object)`.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "====> ref idx: [48 61  0]\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 60/60 [08:26<00:00,  8.44s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "===> valing index: [63, 70, 18, 28]\n",
      "====> ref idx: [6, 43, 33]\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 60/60 [08:25<00:00,  8.43s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "===> valing index: [20 49 55 72]\n",
      "====> ref idx: [61 80 64]\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 60/60 [08:25<00:00,  8.43s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "===> valing index: [ 8 24 32 78]\n",
      "====> ref idx: [62 56 26]\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 60/60 [08:24<00:00,  8.41s/it]\n"
     ]
    }
   ],
   "source": [
    "for i_scene, scene in enumerate(['hotdog','lego','mic','chair']):#'ship','drums','ficus','materials',\n",
    "\n",
    "    cmd = f'--datadir /mnt/new_disk_2/anpei/Dataset/nerf_synthetic/{scene}\\\n",
    "     --dataset_name blender --white_bkgd --imgScale_test {1.0} '\n",
    "\n",
    "    is_finetued = False # set True if rendering with finetuning\n",
    "    if is_finetued:\n",
    "        cmd += f'--ckpt ./runs_fine_tuning/{scene}_1h/ckpts//latest.tar'\n",
    "        pad = 0 #the padding value should be same as your finetuning ckpt\n",
    "    else:\n",
    "        cmd += '--ckpt ./ckpts//mvsnerf-v0.tar'\n",
    "        pad = 24 #the padding value should be same as your finetuning ckpt\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",
    "#     args.use_color_volume = False if not is_finetued else args.use_color_volume\n",
    "\n",
    "    # create models\n",
    "    if i_scene==0 or is_finetued:\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",
    "    args.chunk = 5120\n",
    "    frames = 60\n",
    "\n",
    "\n",
    "    dataset = dataset_dict[args.dataset_name](args, split=datatype)\n",
    "    val_idx = dataset.img_idx\n",
    "    \n",
    "    save_as_image = False\n",
    "    save_dir = f'results/video2'\n",
    "    os.makedirs(save_dir, exist_ok=True)\n",
    "    MVSNet.train()\n",
    "    MVSNet = MVSNet.cuda()\n",
    "    \n",
    "    with torch.no_grad():\n",
    "\n",
    "        if is_finetued:   \n",
    "            # large baselien\n",
    "            imgs_source, proj_mats, near_far_source, pose_source = dataset.read_source_views(device=device)\n",
    "            volume_feature = torch.load(args.ckpt)['volume']['feat_volume']\n",
    "            volume_feature = RefVolume(volume_feature.detach()).cuda()\n",
    "            c2ws_render = nerf_video_path(pose_source['c2ws'].cpu(), N_views=frames)\n",
    "        else:            \n",
    "            # neighboring views with angle distance\n",
    "            c2ws_all = dataset.load_poses_all()\n",
    "            random_selete = torch.randint(0,len(c2ws_all),(1,))     #!!!!!!!!!! you may change this line if rendering a specify view \n",
    "            dis = np.sum(c2ws_all[:,:3,2] * c2ws_all[[random_selete],:3,2], axis=-1)\n",
    "            pair_idx = np.argsort(dis)[::-1][torch.randperm(5)[:3]]\n",
    "            imgs_source, proj_mats, near_far_source, pose_source = dataset.read_source_views(device=device)#pair_idx=pair_idx, \n",
    "            volume_feature, _, _ = MVSNet(imgs_source, proj_mats, near_far_source, pad=pad)\n",
    "            \n",
    "            #####\n",
    "            c2ws_render = gen_render_path(c2ws_all[pair_idx], N_views=frames)\n",
    "            c2ws_render = torch.from_numpy(np.stack(c2ws_render)).float().to(device)\n",
    "\n",
    "            \n",
    "        imgs_source = unpreprocess(imgs_source)\n",
    "        \n",
    "\n",
    "        try:\n",
    "            tqdm._instances.clear() \n",
    "        except Exception:     \n",
    "            pass\n",
    "        \n",
    "        frames = []\n",
    "        img_directions = dataset.directions.to(device)\n",
    "        for i, c2w in enumerate(tqdm(c2ws_render)):\n",
    "            torch.cuda.empty_cache()\n",
    "            \n",
    "            rays_o, rays_d = get_rays(img_directions, c2w)  # both (h*w, 3)\n",
    "            rays = torch.cat([rays_o, rays_d,\n",
    "                     near_far_source[0] * torch.ones_like(rays_o[:, :1]),\n",
    "                     near_far_source[1] * torch.ones_like(rays_o[:, :1])],\n",
    "                    1).to(device)  # (H*W, 3)\n",
    "            \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 = imgs_source.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",
    "            H_crop, W_crop = np.array(rgb_rays.shape[:2])//20\n",
    "#             rgb_rays = rgb_rays[H_crop:-H_crop,W_crop:-W_crop]\n",
    "#             depth_rays_preds = depth_rays_preds[H_crop:-H_crop,W_crop:-W_crop]\n",
    "            img_vis = np.concatenate((rgb_rays*255,depth_rays_preds),axis=1)\n",
    "\n",
    "            frames.append(img_vis.astype('uint8'))\n",
    "#             break\n",
    "    imageio.mimwrite(f'{save_dir}/ft_{scene}_spiral.mp4', np.stack(frames), fps=10, quality=10)\n",
    "# plt.imshow(rgb_rays)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### DTU video rendering"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Found ckpts ['./runs_fine_tuning/dtu_scan1_2h/ckpts//latest.tar']\n",
      "Reloading from ./runs_fine_tuning/dtu_scan1_2h/ckpts//latest.tar\n",
      "==> image down scale: 1.0\n",
      "===> valing index: [32, 24, 23, 44]\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 60/60 [04:15<00:00,  4.25s/it]\n"
     ]
    }
   ],
   "source": [
    "for i_scene, scene in enumerate([1]):# any scene index, like 1,2,3...,,8,21,103,114\n",
    "\n",
    "    cmd = f'--datadir /mnt/data/new_disk/sungx/data/mvs_dataset/DTU/mvs_training/dtu/scan{scene} \\\n",
    "     --dataset_name dtu_ft --imgScale_test {1.0} ' #--use_color_volume\n",
    "    \n",
    "    is_finetued = True # set False if rendering without finetuning\n",
    "    if is_finetued:\n",
    "        cmd += f'--ckpt ./runs_fine_tuning/dtu_scan{scene}_2h/ckpts//latest.tar'\n",
    "    else:\n",
    "        cmd += '--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",
    "    args.use_color_volume = False if not is_finetued else args.use_color_volume\n",
    "\n",
    "\n",
    "    # create models\n",
    "    if i_scene==0 or is_finetued:\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 #the padding value should be same as your finetuning ckpt\n",
    "    args.chunk = 5120\n",
    "    frames = 60\n",
    "\n",
    "\n",
    "    dataset = dataset_dict[args.dataset_name](args, split=datatype)\n",
    "    val_idx = dataset.img_idx\n",
    "    \n",
    "    save_as_image = False\n",
    "    save_dir = f'results/video2'\n",
    "    os.makedirs(save_dir, exist_ok=True)\n",
    "    MVSNet.train()\n",
    "    MVSNet = MVSNet.cuda()\n",
    "    \n",
    "    with torch.no_grad():\n",
    "        \n",
    "        if is_finetued:   \n",
    "            # large baselien\n",
    "            imgs_source, proj_mats, near_far_source, pose_source = dataset.read_source_views(device=device)\n",
    "            volume_feature = torch.load(args.ckpt)['volume']['feat_volume']\n",
    "            volume_feature = RefVolume(volume_feature.detach()).cuda()\n",
    "        else:            \n",
    "            # neighboring views with angle distance\n",
    "            c2ws_all = dataset.load_poses_all()\n",
    "            random_selete = torch.randint(0,len(c2ws_all),(1,)) #!!!!!!!!!! you may change this line if rendering a specify view \n",
    "            dis = np.sum(c2ws_all[:,:3,2] * c2ws_all[[random_selete],:3,2], axis=-1)\n",
    "            pair_idx = np.argsort(dis)[::-1][:3]#[25, 21, 33]#[14,15,24]#\n",
    "            imgs_source, proj_mats, near_far_source, pose_source = dataset.read_source_views(pair_idx=pair_idx, device=device)\n",
    "            volume_feature, _, _ = MVSNet(imgs_source, proj_mats, near_far_source, pad=pad)\n",
    "            \n",
    "        imgs_source = unpreprocess(imgs_source)\n",
    "\n",
    "        c2ws_render = gen_render_path(pose_source['c2ws'].cpu().numpy(), N_views=frames)\n",
    "        c2ws_render = torch.from_numpy(np.stack(c2ws_render)).float().to(device)\n",
    "        \n",
    "        \n",
    "        try:\n",
    "            tqdm._instances.clear() \n",
    "        except Exception:     \n",
    "            pass\n",
    "        \n",
    "        frames = []\n",
    "        img_directions = dataset.directions.to(device)\n",
    "        for i, c2w in enumerate(tqdm(c2ws_render)):\n",
    "            torch.cuda.empty_cache()\n",
    "            \n",
    "            rays_o, rays_d = get_rays(img_directions, c2w)  # both (h*w, 3)\n",
    "            rays = torch.cat([rays_o, rays_d,\n",
    "                     near_far_source[0] * torch.ones_like(rays_o[:, :1]),\n",
    "                     near_far_source[1] * torch.ones_like(rays_o[:, :1])],\n",
    "                    1).to(device)  # (H*W, 3)\n",
    "            \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 = imgs_source.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",
    "#             H_crop, W_crop = np.array(rgb_rays.shape[:2])//20\n",
    "#             rgb_rays = rgb_rays[H_crop:-H_crop,W_crop:-W_crop]\n",
    "#             depth_rays_preds = depth_rays_preds[H_crop:-H_crop,W_crop:-W_crop]\n",
    "            img_vis = np.concatenate((rgb_rays*255,depth_rays_preds),axis=1)\n",
    "            frames.append(img_vis.astype('uint8'))\n",
    "                \n",
    "    imageio.mimwrite(f'{save_dir}/ft_scan{scene}_spiral2.mp4', np.stack(frames), fps=20, quality=10)\n",
    "# plt.imshow(rgb_rays)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# render path generation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "============> rendering dataset <===================\n",
      "34 34 /mnt/new_disk_2/anpei/Dataset/MVSNeRF/nerf_llff_data/flower\n",
      "16.93030140064795\n",
      "===> valing index: [20  6 22  5]\n",
      "====> ref idx: [12 29 11]\n",
      "============> rendering dataset <===================\n",
      "25 25 /mnt/new_disk_2/anpei/Dataset/MVSNeRF/nerf_llff_data/orchids\n",
      "10.866254797560998\n",
      "===> valing index: [12 10 16 19]\n",
      "====> ref idx: [ 8 13 11]\n",
      "============> rendering dataset <===================\n",
      "41 41 /mnt/new_disk_2/anpei/Dataset/MVSNeRF/nerf_llff_data/room\n",
      "8.030018355528686\n",
      "===> valing index: [35 15 38 21]\n",
      "====> ref idx: [14 39 34]\n",
      "============> rendering dataset <===================\n",
      "26 26 /mnt/new_disk_2/anpei/Dataset/MVSNeRF/nerf_llff_data/leaves\n",
      "26.362457265215173\n",
      "===> valing index: [13 11 16  4]\n",
      "====> ref idx: [12 18  8]\n",
      "============> rendering dataset <===================\n",
      "20 20 /mnt/new_disk_2/anpei/Dataset/MVSNeRF/nerf_llff_data/fern\n",
      "12.738972134007064\n",
      "===> valing index: [12 13  5 19]\n",
      "====> ref idx: [17  2  7]\n",
      "============> rendering dataset <===================\n",
      "62 62 /mnt/new_disk_2/anpei/Dataset/MVSNeRF/nerf_llff_data/horns\n",
      "8.315313360692242\n",
      "===> valing index: [33 40 31 59]\n",
      "====> ref idx: [23 32 24]\n",
      "============> rendering dataset <===================\n",
      "55 55 /mnt/new_disk_2/anpei/Dataset/MVSNeRF/nerf_llff_data/trex\n",
      "10.958461140504273\n",
      "===> valing index: [20 21 53 22]\n",
      "====> ref idx: [52 19 47]\n",
      "============> rendering dataset <===================\n",
      "42 42 /mnt/new_disk_2/anpei/Dataset/MVSNeRF/nerf_llff_data/fortress\n",
      "9.564056923447355\n",
      "===> valing index: [21  9 40 25]\n",
      "====> ref idx: [15 20 26]\n",
      "============> rendering dataset <===================\n",
      "34 34 /mnt/new_disk_2/anpei/Dataset/MVSNeRF/nerf_llff_data/flower\n",
      "16.93030140064795\n",
      "===> valing index: [20  6 22  5]\n",
      "====> ref idx: [12 29 11]\n",
      "============> rendering dataset <===================\n",
      "25 25 /mnt/new_disk_2/anpei/Dataset/MVSNeRF/nerf_llff_data/orchids\n",
      "10.866254797560998\n",
      "===> valing index: [12 10 16 19]\n",
      "====> ref idx: [ 8 13 11]\n",
      "============> rendering dataset <===================\n",
      "41 41 /mnt/new_disk_2/anpei/Dataset/MVSNeRF/nerf_llff_data/room\n",
      "8.030018355528686\n",
      "===> valing index: [35 15 38 21]\n",
      "====> ref idx: [14 39 34]\n",
      "============> rendering dataset <===================\n",
      "26 26 /mnt/new_disk_2/anpei/Dataset/MVSNeRF/nerf_llff_data/leaves\n",
      "26.362457265215173\n",
      "===> valing index: [13 11 16  4]\n",
      "====> ref idx: [12 18  8]\n",
      "============> rendering dataset <===================\n",
      "20 20 /mnt/new_disk_2/anpei/Dataset/MVSNeRF/nerf_llff_data/fern\n",
      "12.738972134007064\n",
      "===> valing index: [12 13  5 19]\n",
      "====> ref idx: [17  2  7]\n",
      "============> rendering dataset <===================\n",
      "62 62 /mnt/new_disk_2/anpei/Dataset/MVSNeRF/nerf_llff_data/horns\n",
      "8.315313360692242\n",
      "===> valing index: [33 40 31 59]\n",
      "====> ref idx: [23 32 24]\n",
      "============> rendering dataset <===================\n",
      "55 55 /mnt/new_disk_2/anpei/Dataset/MVSNeRF/nerf_llff_data/trex\n",
      "10.958461140504273\n",
      "===> valing index: [20 21 53 22]\n",
      "====> ref idx: [52 19 47]\n",
      "============> rendering dataset <===================\n",
      "42 42 /mnt/new_disk_2/anpei/Dataset/MVSNeRF/nerf_llff_data/fortress\n",
      "9.564056923447355\n",
      "===> valing index: [21  9 40 25]\n",
      "====> ref idx: [15 20 26]\n",
      "============> rendering dataset <===================\n",
      "===> valing index: [80 86 22 20]\n",
      "====> ref idx: [12 32 44]\n",
      "============> rendering dataset <===================\n",
      "===> valing index: [20 49 55 72]\n",
      "====> ref idx: [61 80 64]\n",
      "============> rendering dataset <===================\n",
      "===> valing index: [ 8 24 32 78]\n",
      "====> ref idx: [62 56 26]\n",
      "============> rendering dataset <===================\n",
      "===> valing index: [63, 70, 18, 28]\n",
      "====> ref idx: [6, 43, 33]\n",
      "============> rendering dataset <===================\n",
      "===> valing index: [79, 74, 91, 68]\n",
      "====> ref idx: [43, 81, 14]\n",
      "============> rendering dataset <===================\n",
      "===> valing index: [38 23  0  5]\n",
      "====> ref idx: [92 69 56]\n",
      "============> rendering dataset <===================\n",
      "===> valing index: [36 63 46 96]\n",
      "====> ref idx: [34 73 94]\n",
      "============> rendering dataset <===================\n",
      "===> valing index: [26 60 13 47]\n",
      "====> ref idx: [48 61  0]\n",
      "============> rendering dataset <===================\n",
      "==> image down scale: 1.0\n",
      "===> valing index: [32, 24, 23, 44]\n"
     ]
    }
   ],
   "source": [
    "render_poses = {}\n",
    "datatype = 'val'\n",
    "for i_scene, scene in enumerate(['xgaze_11images_cropped_colmapCODE']):\n",
    "    # add --use_color_volume if the ckpts are fintuned with this flag\n",
    "    cmd = f'--datadir /home/hengfei/Desktop/research/mvsnerf/xgaze/{scene}  \\\n",
    "     --dataset_name llff --imgScale_test {1.0} \\\n",
    "    --ckpt ./runs_fine_tuning/{scene}-ft/ckpts/latest.tar'\n",
    "\n",
    "    args = config_parser(cmd.split())\n",
    "    args.use_viewdirs = True\n",
    "\n",
    "\n",
    "    print('============> rendering dataset <===================')\n",
    "    dataset = dataset_dict[args.dataset_name](args, split=datatype)\n",
    "    val_idx = dataset.img_idx\n",
    "    \n",
    "\n",
    "    imgs_source, proj_mats, near_far_source, pose_source = dataset.read_source_views(device=device)\n",
    "\n",
    "    c2ws_all = dataset.poses\n",
    "    w2cs, c2ws = pose_source['w2cs'], pose_source['c2ws']\n",
    "    pair_idx = torch.load('configs/pairs.th')[f'{scene}_train']\n",
    "    c2ws_render = get_spiral(c2ws_all[pair_idx], near_far_source, rads_scale = 0.1, N_views=60) \n",
    "    \n",
    "    render_poses[f'{scene}_c2ws'] = c2ws_render\n",
    "    render_poses[f'{scene}_intrinsic'] = pose_source['intrinsics'][0].cpu().numpy()\n",
    "    \n",
    "#     render_poses[f'{scene}_near_far_source'] = near_far_source\n",
    "#     render_poses[f'{scene}_c2ws_no_ft'] = c2ws_render\n",
    "#     render_poses[f'{scene}_intrinsic_no_ft'] = pose_source['intrinsics'][0].cpu().numpy()\n",
    "\n",
    "# for i_scene, scene in enumerate(['flower','orchids', 'room','leaves','fern','horns','trex','fortress']):#'flower','orchids', 'room','leaves','fern','horns','trex','fortress'\n",
    "#     # add --use_color_volume if the ckpts are fintuned with this flag\n",
    "#     cmd = f'--datadir /mnt/new_disk_2/anpei/Dataset/MVSNeRF/nerf_llff_data/{scene}  \\\n",
    "#      --dataset_name llff --imgScale_test {1.0} \\\n",
    "#     --ckpt ./runs_new/runs_fine_tuning/{scene}/ckpts//latest.tar'\n",
    "\n",
    "#     args = config_parser(cmd.split())\n",
    "#     args.use_viewdirs = True\n",
    "\n",
    "\n",
    "#     print('============> rendering dataset <===================')\n",
    "#     dataset = dataset_dict[args.dataset_name](args, split=datatype)\n",
    "#     val_idx = dataset.img_idx\n",
    "    \n",
    "\n",
    "#     imgs_source, proj_mats, near_far_source, pose_source = dataset.read_source_views(device=device)\n",
    "\n",
    "#     c2ws_all = dataset.poses\n",
    "#     w2cs, c2ws = pose_source['w2cs'], pose_source['c2ws']\n",
    "#     pair_idx = torch.load('configs/pairs.th')[f'{scene}_train']\n",
    "#     c2ws_render = get_spiral(c2ws_all[pair_idx], near_far_source, rads_scale = 0.6, N_views=60) \n",
    "    \n",
    "#     render_poses[f'{scene}_c2ws'] = c2ws_render\n",
    "#     render_poses[f'{scene}_intrinsic'] = pose_source['intrinsics'][0].cpu().numpy()\n",
    "# #######################################\n",
    "# for i_scene, scene in enumerate(['ship','mic','chair','lego','drums','ficus','materials','hotdog']):#\n",
    "\n",
    "#     cmd = f'--datadir /mnt/new_disk_2/anpei/Dataset/nerf_synthetic/{scene}  \\\n",
    "#      --dataset_name blender --white_bkgd --imgScale_test {1.0}  \\\n",
    "#     --ckpt /mnt/new_disk_2/anpei/code/MVS-NeRF/runs_fine_tuning/{scene}/ckpts//latest.tar '\n",
    "\n",
    "#     args = config_parser(cmd.split())\n",
    "\n",
    "\n",
    "#     print('============> rendering dataset <===================')\n",
    "#     dataset = dataset_dict[args.dataset_name](args, split=datatype)\n",
    "\n",
    "\n",
    "#     imgs_source, proj_mats, near_far_source, pose_source = dataset.read_source_views(device=device)\n",
    "#     c2ws_render = nerf_video_path(pose_source['c2ws'].cpu(), N_views=60)\n",
    "    \n",
    "#     render_poses[f'{scene}_c2ws'] = c2ws_render.cpu().numpy()\n",
    "#     render_poses[f'{scene}_intrinsic'] = pose_source['intrinsics'][0].cpu().numpy()\n",
    "    \n",
    "# ##################################################\n",
    "# for i_scene, scene in enumerate([1]):\n",
    "\n",
    "#     cmd = f'--datadir /mnt/data/new_disk/sungx/data/mvs_dataset/DTU/mvs_training/dtu/scan{scene}  \\\n",
    "#      --dataset_name dtu_ft --imgScale_test {1.0}   \\\n",
    "#     --ckpt /mnt/new_disk_2/anpei/code/MVS-NeRF/runs_fine_tuning/dtu_scan{scene}/ckpts//latest.tar --netwidth 256 --net_type v0 --use_color_volume'\n",
    "\n",
    "#     args = config_parser(cmd.split())\n",
    "#     args.use_viewdirs = True\n",
    "\n",
    "\n",
    "#     print('============> rendering dataset <===================')\n",
    "#     dataset = dataset_dict[args.dataset_name](args, split=datatype)\n",
    "#     val_idx = dataset.img_idx\n",
    "\n",
    "#     imgs_source, proj_mats, near_far_source, pose_source = dataset.read_source_views(device=device)\n",
    "#     imgs_source = unpreprocess(imgs_source)\n",
    "\n",
    "#     c2ws_render = gen_render_path(pose_source['c2ws'].cpu().numpy(), N_views=60)\n",
    "#     render_poses[f'dtu_c2ws'] = c2ws_render\n",
    "#     render_poses[f'dtu_intrinsic'] = pose_source['intrinsics'][0].cpu().numpy()\n",
    "\n",
    "    \n",
    "torch.save(render_poses, './configs/video_path.th')\n",
    "np.save('./configs/video_path.npy',render_poses)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "rng = np.random.RandomState(234)\n",
    "_EPS = np.finfo(float).eps * 4.0\n",
    "TINY_NUMBER = 1e-6      # float32 only has 7 decimal digits precision\n",
    "\n",
    "def angular_dist_between_2_vectors(vec1, vec2):\n",
    "    vec1_unit = vec1 / (np.linalg.norm(vec1, axis=1, keepdims=True) + TINY_NUMBER)\n",
    "    vec2_unit = vec2 / (np.linalg.norm(vec2, axis=1, keepdims=True) + TINY_NUMBER)\n",
    "    angular_dists = np.arccos(np.clip(np.sum(vec1_unit*vec2_unit, axis=-1), -1.0, 1.0))\n",
    "    return angular_dists\n",
    "\n",
    "\n",
    "def batched_angular_dist_rot_matrix(R1, R2):\n",
    "    '''\n",
    "    calculate the angular distance between two rotation matrices (batched)\n",
    "    :param R1: the first rotation matrix [N, 3, 3]\n",
    "    :param R2: the second rotation matrix [N, 3, 3]\n",
    "    :return: angular distance in radiance [N, ]\n",
    "    '''\n",
    "    assert R1.shape[-1] == 3 and R2.shape[-1] == 3 and R1.shape[-2] == 3 and R2.shape[-2] == 3\n",
    "    return np.arccos(np.clip((np.trace(np.matmul(R2.transpose(0, 2, 1), R1), axis1=1, axis2=2) - 1) / 2.,\n",
    "                             a_min=-1 + TINY_NUMBER, a_max=1 - TINY_NUMBER))\n",
    "\n",
    "\n",
    "def get_nearest_pose_ids(tar_pose, ref_poses, num_select, tar_id=-1, angular_dist_method='vector',\n",
    "                         scene_center=(0, 0, 0)):\n",
    "    '''\n",
    "    Args:\n",
    "        tar_pose: target pose [3, 3]\n",
    "        ref_poses: reference poses [N, 3, 3]\n",
    "        num_select: the number of nearest views to select\n",
    "    Returns: the selected indices\n",
    "    '''\n",
    "    num_cams = len(ref_poses)\n",
    "    # num_select = min(num_select, num_cams-1)\n",
    "    batched_tar_pose = tar_pose[None].repeat(num_cams,axis=0)\n",
    "\n",
    "    if angular_dist_method == 'matrix':\n",
    "        dists = batched_angular_dist_rot_matrix(batched_tar_pose[:, :3, :3], ref_poses[:, :3, :3])\n",
    "    elif angular_dist_method == 'vector':\n",
    "        tar_cam_locs = batched_tar_pose[:, :3, 3]\n",
    "        ref_cam_locs = ref_poses[:, :3, 3]\n",
    "        scene_center = np.array(scene_center)[None, ...]\n",
    "        tar_vectors = tar_cam_locs - scene_center\n",
    "        ref_vectors = ref_cam_locs - scene_center\n",
    "        dists = angular_dist_between_2_vectors(tar_vectors, ref_vectors)\n",
    "    elif angular_dist_method == 'dist':\n",
    "        tar_cam_locs = batched_tar_pose[:, :3, 3]\n",
    "        ref_cam_locs = ref_poses[:, :3, 3]\n",
    "        dists = np.linalg.norm(tar_cam_locs - ref_cam_locs, axis=1)\n",
    "    else:\n",
    "        raise Exception('unknown angular distance calculation method!')\n",
    "\n",
    "    if tar_id >= 0:\n",
    "        assert tar_id < num_cams\n",
    "        dists[tar_id] = 1e3  # make sure not to select the target id itself\n",
    "\n",
    "    sorted_ids = np.argsort(dists)\n",
    "    selected_ids = sorted_ids[:num_select]\n",
    "    # print(angular_dists[selected_ids] * 180 / np.pi)\n",
    "    return selected_ids\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 57,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "============> rendering dataset <===================\n",
      "===> valing index: [20 49 55 72]\n",
      "====> ref idx: [61 80 64]\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "IMAGEIO FFMPEG_WRITER WARNING: input image is not divisible by macro_block_size=16, resizing from (600, 200) to (608, 208) to ensure video compatibility with most codecs and players. To prevent resizing, make your input image divisible by the macro_block_size or set the macro_block_size to 1 (risking incompatibility).\n"
     ]
    }
   ],
   "source": [
    "import glob\n",
    "render_poses = {}\n",
    "datatype = 'val'\n",
    "fourcc = cv2.VideoWriter_fourcc('m', 'p', '4', 'v')\n",
    "# for i_scene, scene in enumerate(['flower']):#\n",
    "#     # add --use_color_volume if the ckpts are fintuned with this flag\n",
    "#     cmd = f'--datadir /mnt/new_disk_2/anpei/Dataset/MVSNeRF/nerf_llff_data/{scene}  \\\n",
    "#      --dataset_name llff --imgScale_test {1.0} \\\n",
    "#     --ckpt ./runs_new/runs_fine_tuning/{scene}/ckpts//latest.tar'\n",
    "\n",
    "#     args = config_parser(cmd.split())\n",
    "#     args.use_viewdirs = True\n",
    "\n",
    "\n",
    "#     print('============> rendering dataset <===================')\n",
    "#     dataset = dataset_dict[args.dataset_name](args, split=datatype)\n",
    "#     val_idx = dataset.img_idx\n",
    "    \n",
    "\n",
    "#     imgs_source, proj_mats, near_far_source, pose_source = dataset.read_source_views(device=device)\n",
    "\n",
    "#     c2ws_all = dataset.poses\n",
    "#     w2cs, c2ws = pose_source['w2cs'], pose_source['c2ws']\n",
    "#     pair_idx = torch.load('configs/pairs.th')[f'{scene}_train']\n",
    "#     c2ws_render = get_spiral(c2ws_all[pair_idx], near_far_source, rads_scale = 0.6, N_views=60) \n",
    "    \n",
    "#     images = []\n",
    "#     for i, c2w in enumerate(c2ws_render):\n",
    "#         nearest_pose_ids = get_nearest_pose_ids(c2w,\n",
    "#                                                 c2ws_all[pair_idx],\n",
    "#                                                 3,\n",
    "#                                                 angular_dist_method='vector')  \n",
    "#         idxs = pair_idx[nearest_pose_ids]\n",
    "        \n",
    "#         im=[]\n",
    "#         List = sorted(glob.glob(f'/mnt/new_disk_2/anpei/Dataset/MVSNeRF/nerf_llff_data/{scene}/images_4/*'))\n",
    "#         for idx in idxs:\n",
    "#             im.append(cv2.resize(cv2.imread(List[idx]),None,fx=0.25,fy=0.25))\n",
    "#         im = np.concatenate(im,axis=1)\n",
    "#         images.append(im[...,::-1])\n",
    "    \n",
    "#     imageio.mimwrite(f'./results/test4/{scene}.mp4', np.stack(images), fps=20, quality=10)\n",
    "\n",
    "# for i_scene, scene in enumerate(['flower','orchids', 'room','leaves','fern','horns','trex','fortress']):#'flower','orchids', 'room','leaves','fern','horns','trex','fortress'\n",
    "#     # add --use_color_volume if the ckpts are fintuned with this flag\n",
    "#     cmd = f'--datadir /mnt/new_disk_2/anpei/Dataset/MVSNeRF/nerf_llff_data/{scene}  \\\n",
    "#      --dataset_name llff --imgScale_test {1.0} \\\n",
    "#     --ckpt ./runs_new/runs_fine_tuning/{scene}/ckpts//latest.tar'\n",
    "\n",
    "#     args = config_parser(cmd.split())\n",
    "#     args.use_viewdirs = True\n",
    "\n",
    "\n",
    "#     print('============> rendering dataset <===================')\n",
    "#     dataset = dataset_dict[args.dataset_name](args, split=datatype)\n",
    "#     val_idx = dataset.img_idx\n",
    "    \n",
    "\n",
    "#     imgs_source, proj_mats, near_far_source, pose_source = dataset.read_source_views(device=device)\n",
    "\n",
    "#     c2ws_all = dataset.poses\n",
    "#     w2cs, c2ws = pose_source['w2cs'], pose_source['c2ws']\n",
    "#     pair_idx = torch.load('configs/pairs.th')[f'{scene}_train']\n",
    "#     c2ws_render = get_spiral(c2ws_all[pair_idx], near_far_source, rads_scale = 0.6, N_views=60) \n",
    "    \n",
    "#     render_poses[f'{scene}_c2ws'] = c2ws_render\n",
    "#     render_poses[f'{scene}_intrinsic'] = pose_source['intrinsics'][0].cpu().numpy()\n",
    "# #######################################\n",
    "for i_scene, scene in enumerate(['mic']):#\n",
    "\n",
    "    cmd = f'--datadir /mnt/new_disk_2/anpei/Dataset/nerf_synthetic/{scene}  \\\n",
    "     --dataset_name blender --white_bkgd --imgScale_test {1.0}  \\\n",
    "    --ckpt /mnt/new_disk_2/anpei/code/MVS-NeRF/runs_fine_tuning/{scene}/ckpts//latest.tar '\n",
    "\n",
    "    args = config_parser(cmd.split())\n",
    "\n",
    "\n",
    "    print('============> rendering dataset <===================')\n",
    "    dataset = dataset_dict[args.dataset_name](args, split=datatype)\n",
    "\n",
    "    c2ws_all = dataset.load_poses_all()\n",
    "    pair_idx = torch.load('configs/pairs.th')[f'{scene}_train']\n",
    "    imgs_source, proj_mats, near_far_source, pose_source = dataset.read_source_views(device=device)\n",
    "    c2ws_render = nerf_video_path(pose_source['c2ws'].cpu(), N_views=60).cpu().numpy()\n",
    "    \n",
    "    \n",
    "    images = []\n",
    "    for i, c2w in enumerate(c2ws_render):\n",
    "        nearest_pose_ids = get_nearest_pose_ids(c2w,\n",
    "                                                c2ws_all[pair_idx],\n",
    "                                                3,\n",
    "                                                angular_dist_method='vector')  \n",
    "        idxs = pair_idx[nearest_pose_ids]\n",
    "        im=[]\n",
    "        List = sorted(glob.glob(f'/mnt/new_disk2/anpei/Dataset/nerf_synthetic/{scene}/train/*'))\n",
    "        for idx in idxs:\n",
    "            temp = cv2.imread(f'/mnt/new_disk2/anpei/Dataset/nerf_synthetic/mic/train/r_{idx}.png',-1)\n",
    "            im.append(cv2.resize(temp,None,fx=0.25,fy=0.25))\n",
    "        im = np.concatenate(im,axis=1)\n",
    "        images.append(im[...,[2,1,0,3]])\n",
    "    \n",
    "    imageio.mimwrite(f'./results/test4/{scene}.mp4', np.stack(images), fps=20, quality=10)\n",
    "    \n",
    "# ##################################################\n",
    "# for i_scene, scene in enumerate([1]):\n",
    "\n",
    "#     cmd = f'--datadir /mnt/data/new_disk/sungx/data/mvs_dataset/DTU/mvs_training/dtu/scan{scene}  \\\n",
    "#      --dataset_name dtu_ft --imgScale_test {1.0}   \\\n",
    "#     --ckpt /mnt/new_disk_2/anpei/code/MVS-NeRF/runs_fine_tuning/dtu_scan{scene}/ckpts//latest.tar --netwidth 256 --net_type v0 --use_color_volume'\n",
    "\n",
    "#     args = config_parser(cmd.split())\n",
    "#     args.use_viewdirs = True\n",
    "\n",
    "\n",
    "#     print('============> rendering dataset <===================')\n",
    "#     dataset = dataset_dict[args.dataset_name](args, split=datatype)\n",
    "#     val_idx = dataset.img_idx\n",
    "\n",
    "#     imgs_source, proj_mats, near_far_source, pose_source = dataset.read_source_views(device=device)\n",
    "#     imgs_source = unpreprocess(imgs_source)\n",
    "\n",
    "#     c2ws_render = gen_render_path(pose_source['c2ws'].cpu().numpy(), N_views=60)\n",
    "#     render_poses[f'dtu_c2ws'] = c2ws_render\n",
    "#     render_poses[f'dtu_intrinsic'] = pose_source['intrinsics'][0].cpu().numpy()\n",
    "\n"
   ]
  }
 ],
 "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
}