{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "0",
   "metadata": {},
   "source": [
    "The notebook demonstrates:\n",
    "1. how to render our data with our blender rendering script.\n",
    "2. basic usage of plibs (eg, to visualize rgbd, point cloud, etc)\n",
    "3. how to save the rendered data into a tar that can be used by our dataloader."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1",
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "\n",
    "import os\n",
    "\n",
    "import numpy as np\n",
    "import open3d as o3d\n",
    "\n",
    "from lito.preprocess_scripts import tar_utils\n",
    "from plibs import data_utils, structures, utils\n",
    "\n",
    "# assume you run the notebook in the notebooks folder\n",
    "repo_root = os.path.abspath(\"..\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2",
   "metadata": {},
   "outputs": [],
   "source": [
    "# download bunny mesh\n",
    "mesh_dir = os.path.join(repo_root, \"assets\", \"mesh\")\n",
    "mesh_filename = os.path.join(mesh_dir, \"bunny.obj\")\n",
    "\n",
    "if not os.path.exists(mesh_filename):\n",
    "    mesh_url = \"http://kunzhou.net/tex-models/bunny.zip\"\n",
    "    os.makedirs(mesh_dir, exist_ok=True)\n",
    "    cmd = f\"cd {mesh_dir}; wget {mesh_url}; unzip bunny.zip\"\n",
    "    os.system(cmd)\n",
    "\n",
    "assert os.path.exists(mesh_filename), f\"{mesh_filename} does not exist\"\n",
    "\n",
    "# mesh_filename = os.path.join(repo_root, \"libraries/blender_rendering/assets/bunny.glb\")\n",
    "\n",
    "# toy settings\n",
    "num_regular_images = 32  # in paper we render 150\n",
    "num_random_images = 5  # in paper we render 100\n",
    "num_cond_images = 2  # num_cond_images = 32\n",
    "num_gen_eval_images = 3  # in paper we render 150\n",
    "\n",
    "\n",
    "assert os.path.exists(mesh_filename), f\"{mesh_filename} does not exist\"\n",
    "\n",
    "# render with blender (which box-normalizes the mesh to [-1, 1])\n",
    "out_root_dir = \"render_data_results\"\n",
    "out_blender_dir = os.path.join(out_root_dir, \"blender_render_results\")\n",
    "out_dict = data_utils.render_multi_mesh_sample(\n",
    "    out_dir=out_blender_dir,\n",
    "    mesh_filenames=[mesh_filename],\n",
    "    min_mesh_scale=1,\n",
    "    max_mesh_scale=1,\n",
    "    # light\n",
    "    light_mode=\"diffuse\",  # \"random\"\n",
    "    # circular camera (no anti-aliasing)\n",
    "    num_regular_images=num_regular_images,\n",
    "    fov=40.0,  # degree\n",
    "    circular_radius=3.5,  # meter\n",
    "    regular_width_px=518,\n",
    "    regular_height_px=518,\n",
    "    # random camera (with anti-aliasing)\n",
    "    num_random_images=num_random_images,\n",
    "    min_fov=20.0,  # degree\n",
    "    max_fov=40.0,  # degree\n",
    "    min_random_radius=1.5,\n",
    "    max_random_radius=3.5,\n",
    "    random_lookat_r=0.25,\n",
    "    random_width_px=518,\n",
    "    random_height_px=518,\n",
    "    # cond camera\n",
    "    num_gen_eval_images=num_gen_eval_images,\n",
    "    num_cond_images=num_cond_images,\n",
    "    cond_width_px=518,\n",
    "    cond_height_px=518,\n",
    "    gen_eval_width_px=518,\n",
    "    gen_eval_height_px=518,\n",
    "    # misc\n",
    "    read_result_to_rgbd=True,\n",
    "    overwrite=True,\n",
    "    max_memory_gb=None,\n",
    "    timeout=None,\n",
    "    blender_device=\"CPU\",  # \"GPU\"\n",
    "    seed=42,\n",
    "    blender_exe=None,\n",
    "    render_gen=True,\n",
    "    rotate_objs=False,\n",
    "    adjust_exposure=True,\n",
    "    exposure_cam_type=\"regular\",\n",
    ")\n",
    "config_dict = out_dict[\"config_dict\"]\n",
    "rgbd_regular: structures.RGBDImage = out_dict[\"rgbd_regular\"]  # (num_frames, num_regular_views, h_regular, w_regular)\n",
    "rgbd_random: structures.RGBDImage = out_dict[\"rgbd_random\"]  # (num_frames, num_random_views, h_random, w_random)\n",
    "rgbd_cond: structures.RGBDImage = out_dict[\"rgbd_cond\"]  # (num_frames, num_cond_views, h_random, w_random)\n",
    "rgbd_gen_eval: structures.RGBDImage = out_dict[\"rgbd_gen_eval\"]  # (num_frames, num_gen_eval_views, h_random, w_random)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3",
   "metadata": {},
   "outputs": [],
   "source": [
    "# (optional) backproject the rgbd images for visualization\n",
    "pcd_regular: structures.PointCloud = rgbd_regular.get_pcd(subsample=4)\n",
    "pcd_random: structures.PointCloud = rgbd_random.get_pcd(subsample=4)\n",
    "pcd_cond: structures.PointCloud = rgbd_cond.get_pcd(subsample=4)\n",
    "pcd_gen_eval: structures.PointCloud = rgbd_gen_eval.get_pcd(subsample=4)\n",
    "\n",
    "# get open3d point cloud\n",
    "o3d_pcd_random: o3d.geometry.PointCloud = pcd_random.get_o3d_pcds()[0]\n",
    "o3d_pcd_regular: o3d.geometry.PointCloud = pcd_regular.get_o3d_pcds()[0]\n",
    "o3d_pcd_cond: o3d.geometry.PointCloud = pcd_cond.get_o3d_pcds()[0]\n",
    "o3d_pcd_gen_eval: o3d.geometry.PointCloud = pcd_gen_eval.get_o3d_pcds()[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4",
   "metadata": {},
   "outputs": [],
   "source": [
    "# (optional) visualize the backprojected points (they should align)\n",
    "world_frame = o3d.geometry.TriangleMesh.create_coordinate_frame(size=1)\n",
    "utils.visualize_mesh_sequence(\n",
    "    [\n",
    "        o3d_pcd_regular,\n",
    "        o3d_pcd_random,\n",
    "        o3d_pcd_cond,\n",
    "        o3d_pcd_gen_eval,\n",
    "    ],\n",
    "    static_meshes=[world_frame],\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5",
   "metadata": {},
   "outputs": [],
   "source": [
    "# save to a local dir\n",
    "data_dir = os.path.join(out_root_dir, \"rgbd_results\")\n",
    "for name, rgbd in [\n",
    "    [\"rgbd_regular\", rgbd_regular],\n",
    "    [\"rgbd_random\", rgbd_random],\n",
    "    [\"rgbd_cond\", rgbd_cond],\n",
    "    [\"rgbd_gen_eval\", rgbd_gen_eval],\n",
    "]:\n",
    "    rgbd.save_as(\n",
    "        out_dir=os.path.join(data_dir, name),\n",
    "        overwrite=True,\n",
    "        mode=\"png\",\n",
    "        background_color=0,\n",
    "        flag_save_space=True,\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6",
   "metadata": {},
   "outputs": [],
   "source": [
    "# save to a local dir to be tarred\n",
    "tar_dir = os.path.join(out_root_dir, \"tar_results\")\n",
    "tar_utils.save_one_sample_from_rendered_rgbds(\n",
    "    data_dir=data_dir,\n",
    "    out_dir=tar_dir,\n",
    "    uid=\"bunny\",\n",
    "    new_uid=\"bunny_r0000\",\n",
    "    num_points=3_000_000,\n",
    "    num_random_views=2,\n",
    "    num_cond_views=2,\n",
    "    num_gen_eval_views=2,\n",
    "    keep_normal=False,\n",
    "    rgbd_save_format=\"png\",\n",
    "    seed=0,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7",
   "metadata": {},
   "outputs": [],
   "source": [
    "# create tar\n",
    "# you can save multiple samples into `tar_results`.\n",
    "# depending on network speed, ideally a tar should be around 1-2 GB\n",
    "\n",
    "tar_id = \"example\"  # uuid.uuid4()\n",
    "tar_name = f\"{tar_id}.tar\"\n",
    "\n",
    "# try to sort by file name (might not work on mac)\n",
    "cmd = f\"cd {tar_dir} && tar --sort=name -cf ../{tar_name} .\"\n",
    "print(cmd, flush=True)\n",
    "return_code = os.system(cmd)\n",
    "\n",
    "if return_code != 0:\n",
    "    cmd = f\"cd {tar_dir} && tar -cf ../{tar_name} .\"\n",
    "    print(cmd, flush=True)\n",
    "    return_code = os.system(cmd)\n",
    "\n",
    "if return_code == 0:\n",
    "    print(f\"Successfully saved {tar_name}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8",
   "metadata": {},
   "outputs": [],
   "source": [
    "# (optionally) save points for tokenizer example\n",
    "st_pcd = rgbd_regular.get_pcd(\n",
    "    subsample=1,\n",
    "    compute_ray_feature=True,\n",
    ")\n",
    "st_pcd = st_pcd.extract_valid_point_cloud(bidx=0)\n",
    "point_xyz_w = st_pcd.xyz_w.squeeze(0)  # (n, 3xyz_w)  [-1, 1]\n",
    "point_rgb = st_pcd.rgb.squeeze(0)  # (n, 3rgb) [0, 1]\n",
    "point_view_dir = st_pcd.captured_view_direction_w.squeeze(0)  # (n, 3xyz_w) normalized, from pinhole to point\n",
    "\n",
    "# slightly shrinks the points to avoid boundary issue for marching cube\n",
    "point_xyz_w = point_xyz_w * 0.9\n",
    "\n",
    "ridxs = np.random.permutation(len(point_xyz_w))[: 2**20]\n",
    "\n",
    "pdict = dict(\n",
    "    point_xyz_w=point_xyz_w[ridxs].cpu().numpy(),  # (n, 3xyz_w)  [-1, 1]\n",
    "    point_rgb=(point_rgb[ridxs] * 255).cpu().numpy().astype(np.uint8),  # (n, 3rgb) [0, 1]\n",
    "    point_view_dir=(((point_view_dir[ridxs] + 1) * 0.5) * 255).cpu().numpy().astype(np.uint8),  # (n, 3xyz_w)\n",
    ")\n",
    "\n",
    "filename = \"assets/bunny.npz\"\n",
    "os.makedirs(os.path.dirname(filename), exist_ok=True)\n",
    "np.savez_compressed(\n",
    "    filename,\n",
    "    **pdict,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "10",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}