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','hog dog')) 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' elif obj_select == 'hog dog': FILE_ID = "11vNlR4lMvV_AVFgVjZmKMrMWGVG7qhNu" MODEL_DIR = 'models' MODEL_NAME = 'diet_nerf_hotdog' 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"""

{caption}

""", 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)