for_seojun.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 파라미터 생성 결과 실험용 코드"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "- 먼저 요놈 돌리쉠"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import glob\n",
    "import json\n",
    "import numpy as np\n",
    "\n",
    "from nerf.provider import *\n",
    "import os\n",
    "import glob\n",
    "import json\n",
    "import tqdm\n",
    "import numpy as np\n",
    "from scipy.spatial.transform import Slerp, Rotation\n",
    "\n",
    "\n",
    "import torch\n",
    "from nerf.network_tcnn import NeRFNetwork\n",
    "from nerf.utils import *\n",
    "from torch.cuda.amp import autocast\n",
    "\n",
    "def get_opt():\n",
    "    import argparse\n",
    "    parser = argparse.ArgumentParser()\n",
    "    parser.add_argument('--path', default='./data/photoshapes/shape09096_rank00',type=str)\n",
    "    parser.add_argument('-O', action='store_true', help=\"equals --fp16 --cuda_ray --preload\")\n",
    "    parser.add_argument('--test', action='store_true', help=\"test mode\")\n",
    "    parser.add_argument('--workspace', type=str, default='workspace')\n",
    "    parser.add_argument('--seed', type=int, default=0)\n",
    "    ### training options\n",
    "    parser.add_argument('--iters', type=int, default=15000, help=\"training iters\")\n",
    "    parser.add_argument('--lr', type=float, default=1e-2, help=\"initial learning rate\")\n",
    "    parser.add_argument('--ckpt', type=str, default='best')\n",
    "    parser.add_argument('--num_rays', type=int, default=4096, help=\"num rays sampled per image for each training step\")\n",
    "    parser.add_argument('--cuda_ray', action='store_true', help=\"use CUDA raymarching instead of pytorch\")\n",
    "    parser.add_argument('--max_steps', type=int, default=1024, help=\"max num steps sampled per ray (only valid when using --cuda_ray)\")\n",
    "    parser.add_argument('--num_steps', type=int, default=512, help=\"num steps sampled per ray (only valid when NOT using --cuda_ray)\")\n",
    "    parser.add_argument('--upsample_steps', type=int, default=0, help=\"num steps up-sampled per ray (only valid when NOT using --cuda_ray)\")\n",
    "    parser.add_argument('--update_extra_interval', type=int, default=16, help=\"iter interval to update extra status (only valid when using --cuda_ray)\")\n",
    "    parser.add_argument('--max_ray_batch', type=int, default=4096, help=\"batch size of rays at inference to avoid OOM (only valid when NOT using --cuda_ray)\")\n",
    "    parser.add_argument('--patch_size', type=int, default=1, help=\"[experimental] render patches in training, so as to apply LPIPS loss. 1 means disabled, use [64, 32, 16] to enable\")\n",
    "    ### network backbone options\n",
    "    parser.add_argument('--fp16', action='store_true', help=\"use amp mixed precision training\")\n",
    "    parser.add_argument('--ff', action='store_true', help=\"use fully-fused MLP\")\n",
    "    parser.add_argument('--tcnn', action='store_true', help=\"use TCNN backend\")\n",
    "    ### dataset options\n",
    "    parser.add_argument('--color_space', type=str, default='srgb', help=\"Color space, supports (linear, srgb)\")\n",
    "    parser.add_argument('--preload', action='store_true', help=\"preload all data into GPU, accelerate training but use more GPU memory\")\n",
    "    # (the default value is for the fox dataset)\n",
    "    parser.add_argument('--bound', type=float, default=2, help=\"assume the scene is bounded in box[-bound, bound]^3, if > 1, will invoke adaptive ray marching.\")\n",
    "    parser.add_argument('--scale', type=float, default=0.33, help=\"scale camera location into box[-bound, bound]^3\")\n",
    "    parser.add_argument('--offset', type=float, nargs='*', default=[0, 0, 0], help=\"offset of camera location\")\n",
    "    parser.add_argument('--dt_gamma', type=float, default=1/128, help=\"dt_gamma (>=0) for adaptive ray marching. set to 0 to disable, >0 to accelerate rendering (but usually with worse quality)\")\n",
    "    parser.add_argument('--min_near', type=float, default=0.2, help=\"minimum near distance for camera\")\n",
    "    parser.add_argument('--density_thresh', type=float, default=10, help=\"threshold for density grid to be occupied\")\n",
    "    parser.add_argument('--bg_radius', type=float, default=-1, help=\"if positive, use a background model at sphere(bg_radius)\")\n",
    "    ### GUI options\n",
    "    parser.add_argument('--gui', action='store_true', help=\"start a GUI\")\n",
    "    parser.add_argument('--W', type=int, default=1920, help=\"GUI width\")\n",
    "    parser.add_argument('--H', type=int, default=1080, help=\"GUI height\")\n",
    "    parser.add_argument('--radius', type=float, default=5, help=\"default GUI camera radius from center\")\n",
    "    parser.add_argument('--fovy', type=float, default=50, help=\"default GUI camera fovy\")\n",
    "    parser.add_argument('--max_spp', type=int, default=64, help=\"GUI rendering max sample per pixel\")\n",
    "    ### experimental\n",
    "    parser.add_argument('--error_map', action='store_true', help=\"use error map to sample rays\")\n",
    "    parser.add_argument('--clip_text', type=str, default='', help=\"text input for CLIP guidance\")\n",
    "    parser.add_argument('--rand_pose', type=int, default=-1, help=\"<0 uses no rand pose, =0 only uses rand pose, >0 sample one rand pose every $ known poses\")\n",
    "\n",
    "\n",
    "\n",
    "    return parser.parse_args(args=[])\n",
    "opt=get_opt()\n",
    "\n",
    "opt.fp16 = True\n",
    "opt.cuda_ray = True\n",
    "opt.preload = True\n",
    "\n",
    "opt.test=True\n",
    "\n",
    "if opt.patch_size > 1:\n",
    "    opt.error_map = False # do not use error_map if use patch-based training\n",
    "    # assert opt.patch_size > 16, \"patch_size should > 16 to run LPIPS loss.\"\n",
    "    assert opt.num_rays % (opt.patch_size ** 2) == 0, \"patch_size ** 2 should be dividable by num_rays.\""
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Network Define\n",
    "- 이놈도 돌리쉠"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
    "model = NeRFNetwork(\n",
    "        encoding=\"hashgrid\",\n",
    "        bound=opt.bound,\n",
    "        cuda_ray=opt.cuda_ray,\n",
    "        density_scale=1,\n",
    "        min_near=opt.min_near,\n",
    "        density_thresh=opt.density_thresh,\n",
    "        bg_radius=opt.bg_radius,\n",
    "    ).to(device)\n",
    "\n",
    "# ckpt_path를 넣으세요\n",
    "ckpt_path='./data/photoshape_weight/shape09096_rank01/checkpoints/ngp.pth'\n",
    "#ckpt_path='hi.pth'\n",
    "# 저장경로\n",
    "opt.workspace='workspace/result'\n",
    "\n",
    "if not os.path.isfile(opt.workspace):\n",
    "    os.makedirs(opt.workspace,exist_ok=True)\n",
    "\n",
    "torch.cuda.empty_cache()\n",
    "ckpt=torch.load(ckpt_path)\n",
    "\n",
    "model.load_state_dict(ckpt['model'],strict=False)\n",
    "model.to(device)\n",
    "model.eval()\n",
    "\n",
    "\n",
    "opt.path='./data/photoshapes/shape09096_rank01'\n",
    "\n",
    "# camera parametre defining!\n",
    "focal=245/(2*np.tan(np.radians(60)/2))\n",
    "intrinsics=np.array([focal,focal,256//2,256//2])\n",
    "\n",
    "H,W=256,256\n",
    "\n",
    "generation_frame=40 # 생성할 뷰 갯수란다.\n",
    "radius=0.5 # 해당 파라미터를 조정하면 줌업을 할 수 있단다"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Line Pose Generation\n",
    "- upper hemisphere interpolation 비디오 생성"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "poses=line_poses(generation_frame,device,radius=radius)\n",
    "data=get_rays(poses,intrinsics,H,W)\n",
    "\n",
    "with autocast():\n",
    "    with torch.no_grad():\n",
    "        result=model.render(data['rays_o'].float(),data['rays_d'].float())\n",
    "        \n",
    "\n",
    "pred=result['image'].reshape(-1,256,256,3)\n",
    "pred=pred.detach().cpu().numpy()\n",
    "#pred=linear_to_srgb(pred).detach().cpu().numpy()\n",
    "#pred=srgb_to_linear(pred).detach().cpu().numpy()\n",
    "pred=(pred*255).astype(np.uint8)\n",
    "\n",
    "#cv2.imwrite(os.path.join(opt.workspace,'sample0.png'), cv2.cvtColor(pred[3], cv2.COLOR_RGB2BGR))\n",
    "imageio.mimwrite(os.path.join(opt.workspace,'sample.mp4'),pred, fps=8, quality=8, macro_block_size=1)\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Random Pose Generation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "poses=rand_poses(generation_frame,device,theta_range=[np.pi/3,np.pi/3],radius=radius)\n",
    "data=get_rays(poses,intrinsics,H,W)\n",
    "with autocast():\n",
    "    with torch.no_grad():\n",
    "        result=model.render(data['rays_o'].float(),data['rays_d'].float())\n",
    "\n",
    "pred=result['image'].reshape(-1,256,256,3)\n",
    "pred=srgb_to_linear(pred).detach().cpu().numpy()\n",
    "pred=(pred*255).astype(np.uint8)\n",
    "\n",
    "#cv2.imwrite(os.path.join(opt.workspace,'sample1.png'), cv2.cvtColor(pred[3], cv2.COLOR_RGB2BGR))\n",
    "imageio.mimwrite(os.path.join(opt.workspace,'sample.mp4'),pred, fps=2, quality=8, macro_block_size=1)\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 훈련에 사용한 transform_test.json 파일이 있을경우\n",
    "- 데이터셋에서 제공하는 camera parameter로 테스트하는 경우란다."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 데이터셋 디렉터리 경로\n",
    "opt.path='./data/photoshapes/shape09096_rank01'\n",
    "\n",
    "metrics = [PSNRMeter(), LPIPSMeter(device=device)]\n",
    "test_loader = NeRFDataset(opt, device=device, type='test').dataloader()\n",
    "trainer = Trainer('ngp', opt, model, device=device, workspace=opt.workspace, fp16=opt.fp16, metrics=metrics, use_checkpoint=opt.ckpt)\n",
    "trainer.test(test_loader, write_video=True)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "hyperdiff",
   "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.9.16"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}