# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     name: python3
# ---

# + [markdown] id="view-in-github" colab_type="text"
# <a href="https://colab.research.google.com/github/stevliu/editnerf/blob/master/editnerf.ipynb" target="_parent"><img src="https://colab.research.google.com/assets/colab-badge.svg" alt="Open In Colab"/></a>

# + [markdown] id="TR1l7vA4UO3q"
# ### Instructions
# To run the code below:
#
# - Run all the cells to first load the relevant code, install the dependencies, and fetch pretrained models. 
# - Finally, run the UI. Please read the UI instructions below for further instructions on how to operate the user interface. 
#
#
# Feel free to check out http://editnerf.csail.mit.edu/ for additional results and details!

# + id="AH1l-GWSTsxd"
# First, clone the repo.
# !git clone https://github.com/stevliu/editnerf.git 1>> setup_log.txt
# %cd editnerf
# Then, install the dependencies that Colab doesn't have yet.
# !pip install configargparse 1>> setup_log.txt
# %cd torchsearchsorted 
# !pip install . 1>> setup_log.txt
# %cd ../

# + id="Bx-0NcyHVvFK"
# !bash scripts/setup_models.sh 1>> setup_log.txt
# %cd ui

# + id="F1WurRgVUWH4"
try: # set up path
    import google.colab, sys, torch
    sys.path.append('/content/editnerf')
    if not torch.cuda.is_available():
        print("Change runtime type to include a GPU.")  
except:
    pass

# + [markdown] id="OBAKsUvtV-vC"
# ## To operate the interface:
# In all edits, first pick a view to edit. Once you have constructed your editing example, click the 'execute' button to execute the edit. 
#
# Particular editing examples can be saved or loaded. Please click the text box on the top right of the interface to select an edit to load. 
#
# Please refer to these instructions for how to construct your own editing examples. 
#
# ### Color Edits: 
# Click the "edit color" button, choose a color, and scribble the desired color change. Optionally, you may select the background brush "BG" and scribble where to keep unchanged. 
#
# ### Shape Removal: 
# Click the "remove shape" button, scribble where you would like to remove. We treat the unmodified parts of the image as the background to keep unchanged. This can be overriden by scribbling on the canvas with the "BG" brush. 
#
# ### Color/Shape Transfer: 
# Click the "color/shape transfer" button, and choose a instance on the bottom to transfer color/shape from. 
#
# ### Shape Addition:
# Click the "add shape" button, and select a instance to copy from on the bottommost panel. A copy interface will show up, and scribble over a region to copy from. Finally, click on the region in the editing canvas the region to paste onto. 
#
#
#

# + id="qAP1aheQUaHr"
import editingapp
from ui_utils import show

# This specifies which dataset to use. 
# The choices are: 'photoshapes/config.txt', 'cars/config.txt', 'dosovitskiy_chairs/config.txt'
config_file = 'photoshapes/config.txt'

# This specifies which object instance to edit. 
# You can preview which instance number corresponds to which instance using the next cell. 
# For PhotoShapes, we find instances 0, 1, 5, 7, and 11 are nice to edit. 
instance_num = 5

# This specifies the number of views to display. 
# Consider reducing this number if the interface is taking too long to load/update.
num_canvases = 3

interface = editingapp.NeRFEditingApp(instance=instance_num, config=config_file, num_canvases=num_canvases)
show(interface)

# + id="FnuKXL4YWVKs"
### This is a util for previewing what the corresponding instance to an instance number looks like. 
import os
from ui_utils import show
from PIL import Image

# This is the dataset. Choices are 'photoshapes', 'dosovitskiy_chairs', 'cars'
dataset = 'photoshapes' 

# This is your instance number. 
# This visualization supports numbers in {0, ..., 20}
instance_num = 5

# Visualize an example view of the instance
instance_path = os.path.join(dataset, 'instances', '{:03d}.png'.format(instance_num))
show(Image.open(instance_path))

StarcoderdataJupyterScriptsDedupFiltered