Spaces:
Build error
Build error
File size: 5,208 Bytes
19677a1 0d35ba8 19677a1 e4b8bbd baad652 19677a1 0d35ba8 0071656 19677a1 0071656 13bc063 0d35ba8 e4b8bbd a0a54d5 0071656 bd597e9 baad652 a0a54d5 b9ba680 bd597e9 19677a1 fd209d1 b19a277 fd209d1 b9ba680 fd209d1 7a6388a b9ba680 0d35ba8 b9ba680 a0a54d5 b9ba680 19677a1 b9ba680 a0a54d5 19677a1 b9ba680 0d35ba8 19677a1 e4b8bbd 6cbae78 074dc08 6cbae78 5272de4 4f54252 c1f7cd5 4f54252 5272de4 0071656 13bc063 1b312d9 519c3dc b9ba680 cf89f7f b9ba680 cf89f7f b9ba680 cf89f7f b9ba680 519c3dc b19a277 519c3dc |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 |
import os
import builtins
import math
import json
import streamlit as st
import gdown
from demo.src.models import load_trained_model
from demo.src.utils import render_predict_from_pose, predict_to_image
st.set_page_config(page_title="DietNeRF")
with open("config.json") as f:
cfg = json.loads(f.read())
MODEL_DIR = "models"
def select_model():
obj_select = st.selectbox("Select a Scene", ("Mic", "Chair", "Lego", "Ship", "Hotdog"))
DIET_NERF_MODEL_NAME = cfg[obj_select]["DIET_NERF_MODEL_NAME"]
DIET_NERF_FILE_ID = cfg[obj_select]["DIET_NERF_FILE_ID"]
NERF_MODEL_NAME = cfg[obj_select]["NERF_MODEL_NAME"]
NERF_FILE_ID = cfg[obj_select]["NERF_FILE_ID"]
return DIET_NERF_MODEL_NAME, DIET_NERF_FILE_ID, NERF_MODEL_NAME, NERF_FILE_ID
st.title("DietNeRF")
caption = (
"**DietNeRF** achieves SoTA few-shot learning capacity in 3D model reconstruction. "
"Thanks to the 2D supervision by `CLIP (aka. Semantic Consisteny Loss)`, "
"it can render novel and challenging views with `ONLY 8 training images`, "
"**outperforming** [original NeRF](https://www.matthewtancik.com/nerf)!"
)
st.markdown(caption)
st.markdown("")
DIET_NERF_MODEL_NAME, DIET_NERF_FILE_ID, NERF_MODEL_NAME, NERF_FILE_ID = select_model()
@st.cache(show_spinner=False)
def download_diet_nerf_model():
os.makedirs(MODEL_DIR, exist_ok=True)
diet_nerf_model_path = os.path.join(MODEL_DIR, DIET_NERF_MODEL_NAME)
url = f"https://drive.google.com/uc?id={DIET_NERF_FILE_ID}"
gdown.download(url, diet_nerf_model_path, quiet=False)
print(f"Model downloaded from google drive: {diet_nerf_model_path}")
def download_nerf_model():
nerf_model_path = os.path.join(MODEL_DIR, NERF_MODEL_NAME)
url = f"https://drive.google.com/uc?id={NERF_FILE_ID}"
gdown.download(url, nerf_model_path, quiet=False)
print(f"Model downloaded from google drive: {nerf_model_path}")
@st.cache(show_spinner=False, allow_output_mutation=True)
def fetch_diet_nerf_model():
model, state = load_trained_model(MODEL_DIR, DIET_NERF_MODEL_NAME)
return model, state
@st.cache(show_spinner=False, allow_output_mutation=True)
def fetch_nerf_model():
model, state = load_trained_model(MODEL_DIR, NERF_MODEL_NAME)
return model, state
diet_nerf_model_path = os.path.join(MODEL_DIR, DIET_NERF_MODEL_NAME)
if not os.path.isfile(diet_nerf_model_path):
download_diet_nerf_model()
nerf_model_path = os.path.join(MODEL_DIR, NERF_MODEL_NAME)
if not os.path.isfile(nerf_model_path):
download_nerf_model()
diet_nerf_model, diet_nerf_state = fetch_diet_nerf_model()
nerf_model, nerf_state = fetch_nerf_model()
pi = math.pi
st.sidebar.markdown(
"""
<style>
.aligncenter {
text-align: center;
}
</style>
<p class="aligncenter">
<img src="https://user-images.githubusercontent.com/77657524/126361638-4aad58e8-4efb-4fc5-bf78-f53d03799e1e.png" width="420" height="400"/>
</p>
""",
unsafe_allow_html=True,
)
st.sidebar.markdown(
"""
<p style='text-align: center'>
<a href="https://github.com/codestella/putting-nerf-on-a-diet" target="_blank">GitHub</a> | <a href="https://www.notion.so/DietNeRF-Putting-NeRF-on-a-Diet-4aeddae95d054f1d91686f02bdb74745" target="_blank">Project Report</a>
</p>
""",
unsafe_allow_html=True,
)
st.sidebar.header("SELECT YOUR VIEW DIRECTION")
theta = st.sidebar.slider(
"Theta", min_value=-pi, max_value=pi, step=0.5, value=0.0, help="Rotational angle in Horizontal direction"
)
phi = st.sidebar.slider(
"Phi", min_value=0.0, max_value=0.5 * pi, step=0.1, value=1.0, help="Rotational angle in Vertical direction"
)
radius = st.sidebar.slider(
"Radius", min_value=2.0, max_value=6.0, step=1.0, value=3.0, help="Distance between object and the viewer"
)
st.markdown("")
with st.spinner("Rendering View..."):
with st.spinner("It may take 2-3 mins. So, why don't you look at our report in the meantime if you haven't already"):
dn_pred_color, _ = render_predict_from_pose(diet_nerf_state, theta, phi, radius)
dn_im = predict_to_image(dn_pred_color)
dn_w, _ = dn_im.size
dn_new_w = int(2 * dn_w)
dn_im = dn_im.resize(size=(dn_new_w, dn_new_w))
n_pred_color, _ = render_predict_from_pose(nerf_state, theta, phi, radius)
n_im = predict_to_image(n_pred_color)
n_w, _ = n_im.size
n_new_w = int(2 * n_w)
n_im = n_im.resize(size=(n_new_w, n_new_w))
diet_nerf_col, nerf_col = st.beta_columns([1, 1])
diet_nerf_col.markdown(
"""<h4 style='text-align: center'>DietNeRF</h4>""", unsafe_allow_html=True
)
diet_nerf_col.image(dn_im, use_column_width=True)
nerf_col.markdown(
"""<h4 style='text-align: center'>NeRF</h4>""", unsafe_allow_html=True
)
nerf_col.image(n_im, use_column_width=True)
st.markdown(
"> 📒 NOTE: The views may look similar to you but see the "
"[Experimental Results](https://www.notion.so/DietNeRF-Putting-NeRF-on-a-Diet-4aeddae95d054f1d91686f02bdb74745#0f6bc8f1008d4765b9b4635999626d4b) "
"section in our report to get a detailed comparison of differences between DietNeRF and NeRF."
)
|