edit_LLFF_app.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import os\n",
    "from matplotlib import pyplot as plt\n",
    "import time\n",
    "from PIL import Image\n",
    "from Model.triplanelite import triplane_fea\n",
    "from Dataloader.LLFF import LLFFDataset\n",
    "from Processing.rendering import novel_views_LLFF, creat_video, render_img\n",
    "from Processing.vis import load_settings, calc_query_emb, calc_feature_dist\n",
    "from Processing.editing import get_comb_img, mkdir_ifnoexit, save_img, edit_color_list\n",
    "from Processing.trainer import SimpleSampler"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 0. general settings\n",
    "device = 'cuda:0' if torch.cuda.is_available() else 'cpu'\n",
    "scene = 'flower'\n",
    "load_features = False\n",
    "show_selection = False\n",
    "\n",
    "# 1. dataset settings\n",
    "datadir= os.path.join('./Dataset/nerf_llff_data', scene)\n",
    "fea_dir = None\n",
    "LLFF_training = LLFFDataset(datadir, fea_path = fea_dir, split='train',load_features=load_features, downsample=8)\n",
    "LLFF_test = LLFFDataset(datadir, fea_path = fea_dir, split='test',load_features=load_features, downsample=8)\n",
    "trainingSampler = SimpleSampler(len(LLFF_training), 4096)\n",
    "downsample = 8\n",
    "\n",
    "# 2. model settings\n",
    "pretrained_model_path = './pre_trained_models/' + scene + '.pth'\n",
    "aabb = torch.tensor([-1.7, 1.7])\n",
    "nerf_model = triplane_fea(aabb = aabb)\n",
    "nerf_model.load_state_dict(torch.load(pretrained_model_path, map_location=torch.device(device)))\n",
    "nerf_model.to(device)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 1. Query a patch in rendering view"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# render a frame from the test dataset\n",
    "rgb_flower, emb_flower, depth_flower, mask_flower = render_img(nerf_model = nerf_model, device= device, Dataset = LLFF_test, img_index = 0, hn = 0, hf = 1, nb_bins = 96, req_others=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# select a patch from the rendered image and calculate the feature distance\n",
    "# modify from N3F official code, https://github.com/dichotomies/N3F\n",
    "settings = load_settings()[scene]\n",
    "factor = downsample\n",
    "r, c = settings['rc']\n",
    "extent = settings['sz']\n",
    "r = int(r * 8 / factor)\n",
    "c = int(c * 8 / factor)\n",
    "extent = int(extent * 8 / factor)\n",
    "img_w, img_h = LLFF_test.img_wh\n",
    "embq, dir_q = calc_query_emb(emb_flower, r, c, extent, rgb=rgb_flower, vis = True)\n",
    "dist = calc_feature_dist(embq, emb_flower)\n",
    "plt.figure(figsize=(4,3))\n",
    "plt.hist(dist.view(-1).cpu().numpy(), bins=20, density=True, alpha=0.5, label='Ditilled Triplanes') \n",
    "plt.show()\n",
    "plt.close()   "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# show example: use the threshold dis_thr to get the mask of the selected region\n",
    "# modify from N3F official code, https://github.com/dichotomies/N3F\n",
    "rgb_j_fg, emb_j_fg, depth_j_fg, mask_j_fg = render_img(nerf_model = nerf_model, device= device, \\\n",
    "                            Dataset = LLFF_test, img_index = 0, hn = 0, hf = 1, nb_bins = 96,\\\n",
    "                            req_others = True, embq=embq, dis_thr=settings['thr'] + settings['margin'], \n",
    "                            foreground=False, show_selection=True)\n",
    "plt.figure(figsize=(20, 8))\n",
    "plt.subplot(1, 3, 1)\n",
    "plt.imshow(rgb_j_fg)\n",
    "plt.subplot(1, 3, 2)\n",
    "plt.imshow(depth_j_fg, cmap = 'gray')\n",
    "plt.subplot(1, 3, 3)\n",
    "plt.imshow(mask_j_fg, cmap = 'gray')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# if you would like to see selection rendering results (videos) you can set show_selection to True:\n",
    "if show_selection:\n",
    "    foldername = time.strftime(\"%Y%m%d-%H%M%S\")\n",
    "    if not os.path.exists('./render_results'):\n",
    "        os.mkdir('./render_results')\n",
    "    folderpath = os.path.join('./render_results', foldername)\n",
    "    os.mkdir(folderpath)\n",
    "    \n",
    "    name = 'wo_selection'\n",
    "    novel_views_path = novel_views_LLFF(folderpath, name, nerf_model, device, LLFF_test, hn = 0, hf = 1, nb_bins = 96,\n",
    "                     req_others = True)\n",
    "    creat_video(novel_views_path, folderpath, name, req_others=False)\n",
    "    name = 'w_selection'\n",
    "    novel_views_path = novel_views_LLFF(folderpath, name, nerf_model, device, LLFF_test, hn = 0, hf = 1, nb_bins = 96,\n",
    "                     req_others = True, dis_thr = settings['thr'] + settings['margin'], \n",
    "                     embq = embq, dist_less=False, show_selection = True)\n",
    "    creat_video(novel_views_path, folderpath, name, req_others=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 2. Save and edit 3D-aware image context"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "save_data = False # make it true if you would like to create your own edited image\n",
    "save_data_path = './Dataset/edited_imgs/' + scene"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if save_data == True:\n",
    "    gt_4, rgb_4, depth_4, mask_4, ray_4 = get_comb_img(nerf_model, LLFF_training, device)\n",
    "    gt_4_s, rgb_4_s, depth_4_s, mask_4_s, ray_4_s = get_comb_img(nerf_model, LLFF_training, device,\n",
    "                                                        embq = embq, dis_thr = settings['thr'] + settings['margin'], show_selection = True)\n",
    "    \n",
    "    # gt_4_s, rgb_4_s, depth_4_s, mask_4_s\n",
    "    save_img(os.path.join(save_data_path, 'gt_4_s.png'), gt_4_s)\n",
    "    save_img(os.path.join(save_data_path, 'rgb_4_s.png'), rgb_4_s)\n",
    "    save_img(os.path.join(save_data_path, 'depth_4_s.png'), depth_4_s)\n",
    "    save_img(os.path.join(save_data_path, 'mask_4_s.png'), mask_4_s)\n",
    "    # gt_4, rgb_4, depth_4\n",
    "    save_img(os.path.join(save_data_path, 'gt_4.png'), gt_4)\n",
    "    save_img(os.path.join(save_data_path, 'rgb_4.png'), rgb_4)\n",
    "    save_img(os.path.join(save_data_path, 'depth_4.png'), depth_4)\n",
    "    # save ray\n",
    "    torch.save(ray_4_s, os.path.join(save_data_path, 'ray_4_s.pt'))\n",
    "    torch.save(ray_4, os.path.join(save_data_path, 'ray_4.pt'))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Then, you can use any image editing tool to edit **rgb_4.png** and save edited image to a folder. Load the edited image as following."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# show one example of the edited image\n",
    "example_edited_rgb = Image.open(os.path.join(save_data_path, 'color_changed','rgb_4_highblue.png'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "example_edited_rgb"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 3. Edit Appearance "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# create a folder to save editing results\n",
    "foldername = time.strftime(\"%Y%m%d-%H%M%S\")\n",
    "savepath_father = os.path.join(\"./editing_results\", foldername)\n",
    "mkdir_ifnoexit(\"./editing_results\")\n",
    "mkdir_ifnoexit(savepath_father)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# setting\n",
    "# --- example path\n",
    "edit_datapath_mlp1 = os.path.join(save_data_path, 'color_changed')\n",
    "edit_datapath_mlp2 = os.path.join(save_data_path, 'others')\n",
    "isf2c_noc2c = True # if True, we use 36KB setting phi_edit; False, we use 4KB setting phi_edit_{c2c}\n",
    "edit_one_image = 'rgb_4_highblue.png' # set None, if you would like edit each image in the folder"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "c_model_dict1 = edit_color_list(edit_datapath_mlp1, save_data_path=save_data_path, nerf_model=nerf_model, \n",
    "                                savepath_father=savepath_father, device=device,\n",
    "                                LLFF_test=LLFF_test, settings=settings, embq=embq,\n",
    "                                isf2c_noc2c = isf2c_noc2c, edit_one_image = edit_one_image)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 4. Layered Editing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# create a folder to save editing results\n",
    "foldername = time.strftime(\"%Y%m%d-%H%M%S\")\n",
    "savepath_father = os.path.join(\"./editing_results\", foldername)\n",
    "mkdir_ifnoexit(\"./editing_results\")\n",
    "mkdir_ifnoexit(savepath_father)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# setting\n",
    "# --- example path\n",
    "edit_datapath_mlp1 = os.path.join(save_data_path, 'color_changed')\n",
    "edit_datapath_mlp2 = os.path.join(save_data_path, 'others')\n",
    "isf2c_noc2c = True # if True, we use 36KB setting phi_edit; False, we use 4KB setting phi_edit_{c2c}\n",
    "# show example of the edited image\n",
    "edit_layer1_image = 'rgb_4_orange.png'\n",
    "edit_layer2_image = 'rgb_4_tone.png'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "c_model_dict1 = edit_color_list(edit_datapath_mlp1, save_data_path=save_data_path, nerf_model=nerf_model, \n",
    "                                savepath_father=savepath_father, device=device,\n",
    "                                LLFF_test=LLFF_test, settings=settings, embq=embq,\n",
    "                                isf2c_noc2c = isf2c_noc2c, edit_one_image = edit_layer1_image)\n",
    "c_model_dict2 = edit_color_list(edit_datapath_mlp2, save_data_path=save_data_path, nerf_model=nerf_model, \n",
    "                                savepath_father=savepath_father, device=device,\n",
    "                                LLFF_test=LLFF_test, settings=settings, embq=embq,\n",
    "                                isf2c_noc2c = isf2c_noc2c, edit_one_image = edit_layer2_image)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for name1, c_model1 in c_model_dict1.items():\n",
    "    for name2, c_model2 in c_model_dict2.items():\n",
    "        mix_model = [c_model1, c_model2]\n",
    "        name = name1 + '_+_' + name2\n",
    "        print(name)\n",
    "        if isf2c_noc2c == True:\n",
    "            novel_views_path = novel_views_LLFF(savepath_father, name, nerf_model, device, LLFF_test, hn = 0, hf = 1, nb_bins = 96,\n",
    "                                dis_thr = settings['thr'] + settings['margin'], embq = embq, dist_less=True,req_others = False, f2c_models = mix_model)\n",
    "        else:\n",
    "            novel_views_path = novel_views_LLFF(savepath_father, name, nerf_model, device, LLFF_test, hn = 0, hf = 1, nb_bins = 96,\n",
    "                                dis_thr = settings['thr'] + settings['margin'], embq = embq, dist_less=True,req_others = False, c2c_models = mix_model)\n",
    "        creat_video(novel_views_path, savepath_father, name, req_others=False)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.10.14"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}