Spaces:
Runtime error
Runtime error
File size: 2,893 Bytes
19677a1 0d35ba8 19677a1 0d35ba8 19677a1 cacc6ce 19677a1 0d35ba8 bd597e9 57a50dc 19677a1 0d35ba8 19677a1 0d35ba8 19677a1 1adb71b 19677a1 b26d474 19677a1 591edcc 927629c 591edcc 8bc9869 1a30119 bd597e9 |
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 |
import os
import builtins
import math
import streamlit as st
import gdown
#from google_drive_downloader import GoogleDriveDownloader as gdd
from demo.src.models import load_trained_model
from demo.src.utils import render_predict_from_pose, predict_to_image
#from demo.src.config import MODEL_DIR, MODEL_NAME, FILE_ID
st.set_page_config(page_title="DietNeRF Demo")
def select_model():
obj_select = st.selectbox("Select Object to Render", ('Chair', 'Lego','Ship'))
if obj_select == 'Chair':
FILE_ID = "17dj0pQieo94TozFv-noSBkXebduij1aM"
MODEL_DIR = 'models'
MODEL_NAME = 'diet_nerf_chair'
elif obj_select == 'Lego':
FILE_ID = "1D9I-qIVMPaxuCHfUWPWMHaoLYtAmCjwI"
MODEL_DIR = 'models'
MODEL_NAME = 'diet_nerf_lego'
elif obj_select == 'Ship':
FILE_ID = "14ZeJ86ETQr8dtu6CFoxU-ifvniHKo_Dt"
MODEL_DIR = 'models'
MODEL_NAME = 'diet_nerf_ship'
return MODEL_DIR,MODEL_NAME,FILE_ID
MODEL_DIR,MODEL_NAME,FILE_ID = select_model()
@st.cache
def download_model():
os.makedirs(MODEL_DIR, exist_ok=True)
_model_path = os.path.join(MODEL_DIR, MODEL_NAME)
# gdd.download_file_from_google_drive(file_id=FILE_ID,
# dest_path=_model_path,
# unzip=True)
url = f'https://drive.google.com/uc?id={FILE_ID}'
gdown.download(url, _model_path, quiet=False)
print(f'model downloaded from google drive: {_model_path}')
@st.cache(show_spinner=False, allow_output_mutation=True)
def fetch_model():
model, state = load_trained_model(MODEL_DIR, MODEL_NAME)
return model, state
model_path = os.path.join(MODEL_DIR, MODEL_NAME)
if not os.path.isfile(model_path):
download_model()
model, state = fetch_model()
pi = math.pi
st.sidebar.image("images/diet-nerf.png", width=310)
st.sidebar.header('SELECT YOUR VIEW DIRECTION')
theta = st.sidebar.slider("Theta", min_value=-pi, max_value=pi,
step=0.5, value=0.)
phi = st.sidebar.slider("Phi", min_value=0., max_value=0.5*pi,
step=0.1, value=1.)
radius = st.sidebar.slider("Radius", min_value=2., max_value=6.,
step=1., value=3.)
caption = "Diet-NeRF achieves SoTA few-shot learning capacity in 3D model reconstruction. " \
"Thanks to the 2D supervision by CLIP (aka semantic loss), " \
"it can render novel and challenging views with ONLY 8 training images, " \
"outperforming original NeRF!"
st.markdown(f""" <h4> {caption} </h4> """,
unsafe_allow_html=True)
with st.spinner("Rendering Image (may take 2-3 mins)..."):
pred_color, _ = render_predict_from_pose(state, theta, phi, radius)
im = predict_to_image(pred_color)
w, _ = im.size
new_w = int(2*w)
im = im.resize(size=(new_w, new_w))
st.image(im, use_column_width=True)
|