run_co3d.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Run CO3D Sequence for Shape from Silhouette"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys, os, glob\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "\n",
    "import skimage.io as sio\n",
    "import matplotlib.pyplot as plt\n",
    "import pathlib"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "#os.environ[\"XLA_PYTHON_CLIENT_PREALLOCATE\"] = \"false\"\n",
    "\n",
    "import jax\n",
    "#jax.config.update('jax_platform_name', 'cpu')\n",
    "\n",
    "import jax.numpy as jnp\n",
    "import fm_render\n",
    "\n",
    "render_jit_ray = jax.jit(fm_render.render_func_rays)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# I recommend grabbing the single sequences, as they're reasonably small. \n",
    "dataset_dir = 'teddybear/34_1479_4753/'\n",
    "input_folder = os.path.join(dataset_dir,'images')\n",
    "co3d_seq = os.path.split(dataset_dir.rstrip('/').lstrip('/'))[-1]\n",
    "output_folder = os.path.join('tmp_out',co3d_seq)\n",
    "NUM_MIXTURE = 40\n",
    "shape_scale = 2.1\n",
    "c_scale = 4.0\n",
    "rand_sphere_size = 40\n",
    "cov_scale = 1.5e-2\n",
    "weight_scale = 1.2\n",
    "LR_RATE = 0.1\n",
    "beta2 = 24.36\n",
    "beta3 = 3.14\n",
    "#beta2, beta3 = jnp.array(fm_render.hyperparams)\n",
    "Nepoch = 10\n",
    "batch_size = 50000\n",
    "target_size = 125000//4"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# do it at some canonical size\n",
    "in_files = sorted(glob.glob(os.path.join(input_folder,'*.jpg')) + glob.glob(os.path.join(input_folder,'*.png')))\n",
    "PYo,PXo = sio.imread(in_files[0]).shape[:2]\n",
    "init_scale = np.prod([PYo,PXo])\n",
    "scales = {}\n",
    "for i in range(10):\n",
    "    scale = 2**i\n",
    "    scales[scale] = init_scale/(scale**2)\n",
    "scale_to_use = sorted([(abs(np.log(v/target_size)),k) for k,v in scales.items() ])[0][1]\n",
    "PY,PX = int(round(PYo/scale_to_use)),int(round(PXo/scale_to_use))\n",
    "scale_to_use,PY,PX"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import skimage\n",
    "import skimage.io as sio\n",
    "import skimage.transform as strans\n",
    "# co3d sequences miss some data\n",
    "valid_inputs = []\n",
    "color_images = []\n",
    "file_map = {}\n",
    "\n",
    "for idx,file in enumerate(in_files):\n",
    "    name = pathlib.Path(file).parts[-1]\n",
    "    img = sio.imread(file)\n",
    "    valid_inputs.append(img.sum() != 0)\n",
    "    new_name = 'frame{:06d}.jpg'.format(sum(valid_inputs))\n",
    "    if valid_inputs[-1] == False:\n",
    "        continue\n",
    "    #print(new_name)\n",
    "    file_map[idx] = sum(valid_inputs)\n",
    "    simg = strans.resize(img,(PY,PX))\n",
    "    color_images.append(simg)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df = pd.read_json(os.path.join(*(dataset_dir.split('/')[:-2] + ['frame_annotations.jgz'])),compression={'method':'gzip'})\n",
    "df2 = df[df.sequence_name == int(co3d_seq.replace('_',''))]\n",
    "fls = []\n",
    "pps = []\n",
    "sizes = []\n",
    "assert(len(df2) == len(valid_inputs))\n",
    "for i,row in enumerate(df2.sort_values('frame_number').itertuples()):\n",
    "    fn, imgd, maskd, view = row[2],row[4],row[6],row[7]\n",
    "    if not valid_inputs[i]:\n",
    "        continue\n",
    "    fl = np.array(view['focal_length'])\n",
    "    pp = np.array(view['principal_point'])\n",
    "    sizeA = list(row[4]['size'])\n",
    "\n",
    "    if 'intrinsics_format' in view and view['intrinsics_format'] == 'ndc_isotropic':\n",
    "        half_image_size_wh_orig = np.array(list(reversed(sizeA))) / 2.0\n",
    "        rescale = half_image_size_wh_orig.min()\n",
    "        # principal point and focal length in pixels\n",
    "        principal_point_px = half_image_size_wh_orig - pp * rescale\n",
    "        focal_length_px = fl * rescale\n",
    "    else:\n",
    "        half_image_size_wh_orig = np.array(list(reversed(sizeA))) / 2.0\n",
    "        # principal point and focal length in pixels\n",
    "        principal_point_px = (\n",
    "            -1.0 * (pp - 1.0) * half_image_size_wh_orig\n",
    "        )\n",
    "        focal_length_px = fl * half_image_size_wh_orig\n",
    "\n",
    "    fls.append(focal_length_px)\n",
    "    pps.append(principal_point_px)\n",
    "\n",
    "    sizes.append(sizeA)\n",
    "assert(np.array(sizes).std(0).sum() == 0) # same sizes\n",
    "pp = np.array(pps).mean(0)\n",
    "fl = np.array(fls).mean(0).mean()\n",
    "meanpp = (np.array([pp[1],pp[0]])/np.array(sizes).mean(0)).mean() \n",
    "assert(abs(meanpp - 0.5) < 1e-3) # basically center of frame\n",
    "fl = fl/scale_to_use"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "poses = []\n",
    "depths = []\n",
    "masks = []\n",
    "import skimage.io as sio\n",
    "import skimage.transform as sktrans\n",
    "import transforms3d\n",
    "for i,row in enumerate(df2.sort_values('frame_number').itertuples()):\n",
    "    fn, imgd, maskd, view = row[2],row[4],row[6],row[7]\n",
    "    depthd = row[5]\n",
    "    if not valid_inputs[i]:\n",
    "        continue\n",
    "    #maskd = maskd['path']#[maskd['path'].index(co3d_seq):]\n",
    "    #imgd = imgd['path']#[imgd['path'].index(co3d_seq):]\n",
    "    mask = np.clip(sio.imread(maskd['path'])/253,0,1) #> 0\n",
    "    masks.append(sktrans.resize(mask,(PY,PX),anti_aliasing=True,order=0))\n",
    "    \n",
    "    Rmat = np.array(view['R'])\n",
    "    Tvec = np.array(view['T'])\n",
    "    Tvec = -Rmat @ Tvec\n",
    "    q = transforms3d.quaternions.mat2quat(Rmat.T)\n",
    "    poses.append((q,Tvec))\n",
    "    \n",
    "    depth_r = sio.imread(depthd['path'])#.astype(float)\n",
    "    depth_m = sio.imread(depthd['mask_path']).astype(float)\n",
    "    \n",
    "    depth_r_s = depth_r.shape\n",
    "    depth_r = depthd['scale_adjustment']*np.frombuffer(depth_r,dtype=np.float16).astype(np.float32).reshape(depth_r_s)\n",
    "\n",
    "    valid_d = (depth_r > 0)\n",
    "\n",
    "    depth_r[~valid_d] = np.nan\n",
    "    depth_r[~(depth_m >0)] = np.nan\n",
    "\n",
    "    depth_r = sktrans.resize(depth_r,(PY,PX),anti_aliasing=False,order=0)\n",
    "    depths.append(depth_r)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if os.path.exists(os.path.join(dataset_dir,'pointcloud.ply')):\n",
    "    import trimesh\n",
    "    mesh_tri = trimesh.load(os.path.join(dataset_dir,'pointcloud.ply'))\n",
    "    pt_cld = mesh_tri.vertices\n",
    "    import sklearn.mixture as mixture\n",
    "\n",
    "    idx2 = np.arange(pt_cld.shape[0])\n",
    "    np.random.shuffle(idx2)\n",
    "    clf = mixture.GaussianMixture(40)\n",
    "    clf.fit(pt_cld[idx2[:10000]])\n",
    "\n",
    "    pt_cld_shape_scale = float(pt_cld.std(0).mean())*3\n",
    "    center = pt_cld.mean(0)\n",
    "else:         \n",
    "    pt_cld_shape_scale = 3.0\n",
    "    center = np.zeros(3,dtype=np.float32)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "SCALE_MUL_FACTOR = shape_scale/pt_cld_shape_scale\n",
    "SCALE_MUL_FACTOR"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "height, width = PY,PX\n",
    "cx = (PX-1)/2\n",
    "cy = (PY-1)/2\n",
    "K = np.array([[fl, 0, cx],[0,fl,cy],[0,0,1]])\n",
    "pixel_list = (np.array(np.meshgrid(width-np.arange(width)-1,height-np.arange(height)-1,[0]))[:,:,:,0]).reshape((3,-1)).T\n",
    "camera_rays = (pixel_list - K[:,2])/np.diag(K)\n",
    "camera_rays[:,-1] = 1\n",
    "cameras_list = []\n",
    "for quat,trans in poses:\n",
    "    R = transforms3d.quaternions.quat2mat(quat)\n",
    "    camera_rays2 = camera_rays @ R\n",
    "    t = np.tile(trans[None],(camera_rays2.shape[0],1))\n",
    "    \n",
    "    rays_trans = np.stack([camera_rays2,t],1)\n",
    "    cameras_list.append(rays_trans)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from util import image_grid"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# random init settings\n",
    "rand_mean = center+pt_cld_shape_scale*np.random.multivariate_normal(mean=[0,0,0],cov=cov_scale*np.identity(3),size=NUM_MIXTURE)\n",
    "rand_weight_log = jnp.log(weight_scale*np.ones(NUM_MIXTURE)/NUM_MIXTURE)\n",
    "rand_prec = jnp.array([np.identity(3)*rand_sphere_size/pt_cld_shape_scale for _ in range(NUM_MIXTURE)])\n",
    "rand_color = jnp.array(np.random.randn(NUM_MIXTURE,3))\n",
    "\n",
    "init_alphas = []\n",
    "init_depths = []\n",
    "render_jit = jax.jit(fm_render.render_func_rays)\n",
    "\n",
    "for ray_trans in cameras_list[:36]:\n",
    "    est_depth, est_alpha, est_norm, est_w = render_jit(rand_mean,rand_prec,rand_weight_log,ray_trans,beta2/shape_scale,beta3)\n",
    "\n",
    "    est_depth = np.array(est_depth)\n",
    "    est_depth[est_alpha < 0.5] = np.nan\n",
    "    init_alphas.append(est_alpha.reshape((PY,PX)))\n",
    "    init_depths.append(est_depth.reshape((PY,PX)))\n",
    "\n",
    "image_grid(init_alphas,6,6,rgb=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def objective(params,camera_rays,beta2,beta3,true_alpha,true_color):\n",
    "    CLIP_ALPHA = 1e-7\n",
    "    means,prec,weights_log,colors = params\n",
    "    est_depth, est_alpha, est_norm, est_w = fm_render.render_func_rays(means,prec,weights_log,camera_rays,beta2,beta3)\n",
    "    est_color = est_w.T @ (jnp.tanh(colors)*0.5+0.5)\n",
    "    est_alpha = jnp.clip(est_alpha,CLIP_ALPHA,1-CLIP_ALPHA)\n",
    "    mask_loss = - ((true_alpha * jnp.log(est_alpha)) + (1-true_alpha)*jnp.log(1-est_alpha))\n",
    "    cdiff = jnp.abs( (true_color-est_color)*true_alpha[:,None] )\n",
    "    return mask_loss.mean() + c_scale*cdiff.mean()\n",
    "grad_render3 = jax.value_and_grad(objective)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import optax\n",
    "from tqdm.notebook import tqdm\n",
    "from util import DegradeLR\n",
    "\n",
    "vecM = jnp.array([[1,SCALE_MUL_FACTOR]])[:,:,None]\n",
    "\n",
    "all_rays = jnp.vstack(cameras_list)\n",
    "train_size = all_rays.shape[0]\n",
    "Niter_epoch = int(round(train_size/batch_size))\n",
    "\n",
    "def irc(x): return int(round(x))\n",
    "\n",
    "# babysit learning rates\n",
    "adjust_lr = DegradeLR(LR_RATE,0.5,irc(Niter_epoch*0.4),irc(Niter_epoch*0.1),-1e-4)\n",
    "\n",
    "optimizer = optax.adam(adjust_lr.step_func)\n",
    "\n",
    "tmp = [rand_mean,rand_prec,rand_weight_log,rand_color]\n",
    "\n",
    "opt_state = optimizer.init(tmp)\n",
    "\n",
    "all_sils = jnp.hstack([_.ravel() for _ in masks]).astype(jnp.float32)\n",
    "all_colors = jnp.hstack([_.ravel() for _ in color_images]).astype(jnp.float32).reshape((-1,3))\n",
    "all_colors = all_colors**(1/2.2)\n",
    "\n",
    "losses = []\n",
    "opt_configs = []\n",
    "outer_loop = tqdm(range(Nepoch), desc=\" epoch\", position=0)\n",
    "\n",
    "rand_idx = np.arange(train_size)\n",
    "params = tmp\n",
    "def inner_iter(j_idx,rand_idx_local,opt_state,p):\n",
    "    idx = jax.lax.dynamic_slice(rand_idx_local,[j_idx*batch_size],[batch_size])\n",
    "\n",
    "    val,g = grad_render3([p[0]*SCALE_MUL_FACTOR,p[1]/SCALE_MUL_FACTOR,p[2],p[3]],vecM*all_rays[idx],\n",
    "                         beta2/(shape_scale),beta3,all_sils[idx],all_colors[idx])   \n",
    "    updates, opt_state = optimizer.update(g, opt_state,p)\n",
    "    p = optax.apply_updates(p, updates)\n",
    "    return val, opt_state, p \n",
    "jax_iter = jax.jit(inner_iter)\n",
    "done = False\n",
    "for i in outer_loop:\n",
    "    np.random.shuffle(rand_idx)\n",
    "    rand_idx_jnp = jnp.array(rand_idx)\n",
    "\n",
    "    for j in tqdm(range(Niter_epoch), desc=\" iteration\", position=1, leave=False):\n",
    "        opt_configs.append(list(params))\n",
    "        val,opt_state,params = jax_iter(j,rand_idx_jnp,opt_state,params)\n",
    "        val = float(val)\n",
    "        losses.append(val)\n",
    "\n",
    "        if adjust_lr.add(val):\n",
    "            done = True\n",
    "            break\n",
    "        outer_loop.set_description(\" loss {:.3f}\".format(val))\n",
    "    if done:\n",
    "        break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.plot(losses)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "final_mean, final_prec, final_weight_log,final_color = params"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "result_depths = []\n",
    "result_alphas = []\n",
    "results_colors = []\n",
    "\n",
    "for ray_trans in cameras_list:\n",
    "    est_depth, est_alpha, est_norm, est_w = render_jit(final_mean,final_prec,final_weight_log,ray_trans,beta2/shape_scale,beta3)\n",
    "    est_w = est_w.T\n",
    "    est_w = est_w/jnp.maximum(est_w.sum(axis=1,keepdims=True),1e-7)\n",
    "    est_color = est_alpha[:,None] * np.array(est_w @ (jnp.tanh(final_color)*0.5+0.5))**(2.2)\n",
    "\n",
    "    est_depth = np.array(est_depth)\n",
    "    est_alpha = np.array(est_alpha)\n",
    "    est_depth[est_alpha < 0.5] = np.nan\n",
    "    #est_color[est_alpha < 0.5] = np.nan\n",
    "\n",
    "    result_depths.append(est_depth.reshape((PY,PX)))\n",
    "    result_alphas.append(est_alpha.reshape((PY,PX)))\n",
    "    results_colors.append((est_color).reshape((PY,PX,3)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.subplot(1,3,1)\n",
    "plt.imshow(result_alphas[-1])\n",
    "plt.axis('off')\n",
    "\n",
    "plt.subplot(1,3,2)\n",
    "plt.imshow(result_depths[-1])\n",
    "plt.axis('off')\n",
    "plt.subplot(1,3,3)\n",
    "plt.imshow(est_color.reshape((PY,PX,3)),interpolation='nearest')\n",
    "plt.axis('off')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from scipy.stats import trim_mean\n",
    "errs = []\n",
    "d1f = np.hstack([_.ravel() for _ in  depths]).ravel()\n",
    "d2f = np.hstack([_.ravel() for _ in result_depths]).ravel()\n",
    "\n",
    "mask = (all_sils !=0 ) & (~np.isnan(d1f)) & (~np.isnan(d2f)) & (d1f !=0) \n",
    "\n",
    "trim_mean(abs(d1f[mask]-d2f[mask]),0.1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "image_grid(masks,rows=3,cols=5,rgb=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "max_frame = len(poses)\n",
    "FWD_BCK_TIMES = 4\n",
    "THRESH_IDX = np.where(np.array(losses)/min(losses) < 1.02)[0][0]\n",
    "USE_FIRST_N_FRAC = THRESH_IDX/len(losses)\n",
    "N_FRAMES = max_frame*FWD_BCK_TIMES\n",
    "opt_to_use = np.round(np.linspace(0,int(np.floor(len(opt_configs)*USE_FIRST_N_FRAC-1)),N_FRAMES)).astype(int)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "THRESH_IDX/len(losses)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.plot(losses[:THRESH_IDX])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "frame_idxs = []\n",
    "frame_list = list(range(max_frame))\n",
    "for i in range(FWD_BCK_TIMES):\n",
    "    if (i % 2) == 0:\n",
    "        frame_idxs += frame_list\n",
    "    else:\n",
    "        frame_idxs += frame_list[::-1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "full_res_alpha = []\n",
    "full_res_depth = []\n",
    "full_res_color = []\n",
    "\n",
    "for r_idx,c_idx in zip(frame_idxs,opt_to_use):\n",
    "    p = opt_configs[c_idx]\n",
    "    ray_trans = cameras_list[r_idx]\n",
    "    est_depth, est_alpha, est_norm, est_w = render_jit(p[0],p[1],p[2],ray_trans,beta2/shape_scale,beta3)\n",
    "    est_w = est_w.T\n",
    "    est_w = est_w/jnp.maximum(est_w.sum(axis=1,keepdims=True),1e-7)\n",
    "    est_color = est_alpha[:,None] * np.array(est_w @ (jnp.tanh(p[3])*0.5+0.5))**(2.2)\n",
    "\n",
    "    est_alpha = np.array(est_alpha)\n",
    "    est_depth = np.array(est_depth)\n",
    "    est_depth[est_alpha < 0.5] = np.nan\n",
    "    #est_color[est_alpha < 0.5] = np.nan\n",
    "\n",
    "    full_res_alpha.append(est_alpha.reshape((PY,PX)))\n",
    "    full_res_depth.append(est_depth.reshape((PY,PX)))\n",
    "    full_res_color.append(est_color.reshape((PY,PX,3)))\n",
    "\n",
    "    print('.',end='')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if os.path.exists(output_folder):\n",
    "    import shutil\n",
    "    shutil.rmtree(output_folder)\n",
    "os.makedirs(output_folder, exist_ok=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "vecr = np.hstack([_.ravel() for _ in full_res_depth])\n",
    "vecr = vecr[~np.isnan(vecr)]\n",
    "vmin = np.percentile(vecr,5)\n",
    "vmax = np.percentile(vecr,95)\n",
    "vscale = vmax-vmin"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib\n",
    "from PIL import Image, ImageDraw, ImageFont\n",
    "start_f = 0\n",
    "avg_size = np.array([PX,PY])\n",
    "fsize = irc(96/4)\n",
    "\n",
    "font = ImageFont.truetype('Roboto-Regular.ttf', size=irc(avg_size[0]/8))\n",
    "cmap = matplotlib.cm.get_cmap('viridis')\n",
    "cmap2 = matplotlib.cm.get_cmap('magma')\n",
    "\n",
    "for i,mask_res in enumerate(full_res_alpha):\n",
    "    r_idx = frame_idxs[i]\n",
    "    #img1 = ground_images[r_idx]/255.0*np.clip(full_masks[r_idx] > .1,0.3,1)[:,:,None]\n",
    "    #img2 = ground_images[r_idx]*np.clip((mask_res)**0.4,0.05,1)[:,:,None]\n",
    "    img2 = full_res_color[i]#np.tile(mask_res[:,:,None],(1,1,3))\n",
    "    img_gt_mask = np.tile(masks[r_idx][:,:,None],(1,1,3))\n",
    "\n",
    "    true_alpha = masks[r_idx]\n",
    "\n",
    "    est_alpha = jnp.clip(mask_res,1e-6,1-1e-6)\n",
    "    mask_loss = - ((true_alpha * jnp.log(est_alpha)) + (1-true_alpha)*jnp.log(1-est_alpha))\n",
    "    loss_viz = cmap2(0.25*mask_loss)[:,:,:3]\n",
    "\n",
    "    depth = cmap((full_res_depth[i]-vmin)/vscale)[:,:,:3]\n",
    "    img2 = np.concatenate((color_images[r_idx],img_gt_mask,loss_viz,img2, depth), axis=1)\n",
    "    int_img = np.round(img2*255).astype(np.uint8)\n",
    "    pil_img = Image.fromarray(int_img)\n",
    "    d1 = ImageDraw.Draw(pil_img)\n",
    "    d1.text((avg_size[0]*1.1, irc(fsize*0.1)), \"Iteration: {:3d}\\nEpoch: {:.1f}\".format(opt_to_use[i],opt_to_use[i]/Niter_epoch), ha='center',font=font,fill=(180, 180, 180))\n",
    "    d1.text((avg_size[0]*1.3, irc(avg_size[1]-fsize*1.5)), \"Target Mask\", font=font,fill=(255, 255, 255),ha='center')\n",
    "    d1.text((avg_size[0]*2.4, irc(avg_size[1]-fsize*1.5)), \"Loss\", font=font,fill=(255, 255, 255),ha='center',align='center')\n",
    "    d1.text((avg_size[0]*3.3, irc(avg_size[1]-fsize*2.5)), \"Estimated\\nColor\", font=font,fill=(255, 255, 255),ha='center',align='center')\n",
    "    d1.text((avg_size[0]*4.3, irc(avg_size[1]-fsize*2.5)), \"Estimated\\nDepth\", font=font,fill=(255, 255, 255),ha='center',align='center')\n",
    "\n",
    "    img3 = np.array(pil_img)\n",
    "    \n",
    "    \n",
    "    sio.imsave('{}/{:03d}.jpg'.format(output_folder,i),img3,quality=95)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure(figsize=(18,8))\n",
    "plt.imshow(img3)\n",
    "plt.axis('off')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import subprocess\n",
    "if os.path.exists('{}.mp4'.format(output_folder)):\n",
    "    os.remove('{}.mp4'.format(output_folder))\n",
    "subprocess.call(' '.join(['/usr/local/bin/ffmpeg',\n",
    "                 '-framerate','60',\n",
    "                 '-i','{}/%03d.jpg'.format(output_folder),\n",
    "                 '-vf','\\\"pad=ceil(iw/2)*2:ceil(ih/2)*2\\\"',\n",
    "                 '-c:v','h264',\n",
    "                 '-pix_fmt','yuv420p',\n",
    "                 '{}.mp4'.format(output_folder)]),shell=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# render with compositing for weights and consistent geometry\n",
    "render_jit2 = jax.jit(fm_render.render_func_rays_hffm)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "points_export = []\n",
    "colors_export = []\n",
    "normals_export = []\n",
    "thesh_min = 0.9\n",
    "for i,ray_trans in enumerate(cameras_list):\n",
    "    est_depth, est_alpha, est_norm, est_w = render_jit2(final_mean,final_prec,final_weight_log,ray_trans)\n",
    "\n",
    "    est_color = np.array(est_w.T @ (jnp.tanh(final_color)*0.5+0.5))**(2.2)\n",
    "    # using the actual high detail colors\n",
    "    # instead of the modeled colors, since we only have 40\n",
    "    export_c = np.round(color_images[i]*255).astype(np.uint8).reshape((-1,3))\n",
    "    \n",
    "    est_3d = est_depth[:,None]*ray_trans[:,0]+ray_trans[:,1] \n",
    "    \n",
    "    est_3d = np.array(est_3d)\n",
    "    est_alpha = np.array(est_alpha)\n",
    "    export_cond = (est_alpha > thesh_min) & (est_w.max(axis=0) > thesh_min)\n",
    "\n",
    "    points_export.append(est_3d[export_cond])\n",
    "    colors_export.append(export_c[export_cond])\n",
    "    normals_export.append(est_norm[export_cond])\n",
    "points_export = np.concatenate(points_export)\n",
    "colors_export = np.concatenate(colors_export)\n",
    "normals_export = np.concatenate(normals_export)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import open3d as o3d\n",
    "o3d_cld = o3d.geometry.PointCloud(o3d.cpu.pybind.utility.Vector3dVector(points_export))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "o3d_cld.colors = o3d.cpu.pybind.utility.Vector3dVector(colors_export[:,:3].astype(float)/255.0)\n",
    "o3d_cld.normals = o3d.cpu.pybind.utility.Vector3dVector(normals_export)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "o3d.io.write_point_cloud(\"export_hffm.ply\", o3d_cld)\n",
    "# to get mesh, run\n",
    "#  ./PoissonRecon --in export_hffm.ply --out mesh.ply\n",
    "# using https://github.com/mkazhdan/PoissonRecon"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# confirm that both rendering forms are roughly comparable\n",
    "result_depths2 = []\n",
    "result_alphas2 = []\n",
    "results_colors2 = []\n",
    "\n",
    "\n",
    "for ray_trans in cameras_list:\n",
    "    est_depth, est_alpha, est_norm, est_w = render_jit2(final_mean,final_prec,final_weight_log,ray_trans)\n",
    "    est_w = est_w.T\n",
    "    #est_w = est_w/jnp.maximum(est_w.sum(axis=1,keepdims=True),1e-7)\n",
    "    est_color = est_w @ (jnp.tanh(final_color)*0.5+0.5)\n",
    "    est_color = np.array(est_color)\n",
    "\n",
    "    est_depth = np.array(est_depth)\n",
    "    est_alpha = np.array(est_alpha)\n",
    "    est_depth[est_alpha < thesh_min] = np.nan\n",
    "    est_color[est_alpha < thesh_min] = np.nan\n",
    "\n",
    "    result_depths2.append(est_depth.reshape((PY,PX)))\n",
    "    result_alphas2.append(est_alpha.reshape((PY,PX)))\n",
    "    results_colors2.append(est_color.reshape((PY,PX,3)))\n",
    "    break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.imshow(result_depths2[-1])\n",
    "plt.figure()\n",
    "plt.imshow(result_depths[-1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.imshow(results_colors2[-1])\n",
    "plt.figure()\n",
    "plt.imshow(results_colors[-1])"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}