Code refactoring

app.py

@@ -6,55 +6,52 @@ import json
 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")
 
 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"))
-    MODEL_DIR = "models"
-    MODEL_NAME = cfg[obj_select]["DIET_NERF_MODEL_NAME"]
-    FILE_ID = cfg[obj_select]["FILE_ID"]
-    return MODEL_DIR, MODEL_NAME, FILE_ID
+    DIET_NERF_MODEL_NAME = cfg[obj_select]["DIET_NERF_MODEL_NAME"]
+    DIET_NERF_FILE_ID = cfg[obj_select]["DIET_NERF_FILE_ID"]
+    return DIET_NERF_MODEL_NAME, DIET_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 loss), "
+    "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!"
 )
 st.markdown(caption)
 st.markdown("")
-MODEL_DIR, MODEL_NAME, FILE_ID = select_model()
+DIET_NERF_MODEL_NAME, DIET_NERF_FILE_ID = select_model()
 
 
 @st.cache(show_spinner=False)
 def download_model():
     os.makedirs(MODEL_DIR, exist_ok=True)
-    _model_path = os.path.join(MODEL_DIR, MODEL_NAME)
-    url = f"https://drive.google.com/uc?id={FILE_ID}"
+    _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, _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)
+    model, state = load_trained_model(MODEL_DIR, DIET_NERF_MODEL_NAME)
     return model, state
 
 
-model_path = os.path.join(MODEL_DIR, MODEL_NAME)
+model_path = os.path.join(MODEL_DIR, DIET_NERF_MODEL_NAME)
 if not os.path.isfile(model_path):
     download_model()
 

config.json

@@ -1,26 +1,31 @@
 {
     "Chair": {
-        "FILE_ID": "17dj0pQieo94TozFv-noSBkXebduij1aM",
+        "DIET_NERF_FILE_ID": "17dj0pQieo94TozFv-noSBkXebduij1aM",
+        "NERF_FILE_ID": "17dj0pQieo94TozFv-noSBkXebduij1aM",
         "DIET_NERF_MODEL_NAME": "diet_nerf_chair",
         "NERF_MODEL_NAME": "nerf_chair"
     },
     "Lego": {
-        "FILE_ID": "1UaUKgShJTKp5zNQdXvVhkyxgU-2YU3Ay",
+        "DIET_NERF_FILE_ID": "1UaUKgShJTKp5zNQdXvVhkyxgU-2YU3Ay",
+        "NERF_FILE_ID": "17dj0pQieo94TozFv-noSBkXebduij1aM",
         "DIET_NERF_MODEL_NAME": "diet_nerf_lego",
         "NERF_MODEL_NAME": "nerf_lego"
     },
     "Ship": {
-        "FILE_ID": "1Eb4eZ9vn6FcbqLIaBoW-6NlyWr-y0eTW",
+        "DIET_NERF_FILE_ID": "1Eb4eZ9vn6FcbqLIaBoW-6NlyWr-y0eTW",
+        "NERF_FILE_ID": "17dj0pQieo94TozFv-noSBkXebduij1aM",
         "DIET_NERF_MODEL_NAME": "diet_nerf_ship",
         "NERF_MODEL_NAME": "nerf_ship"
     },
     "Hotdog": {
-        "FILE_ID": "11vNlR4lMvV_AVFgVjZmKMrMWGVG7qhNu",
+        "DIET_NERF_FILE_ID": "11vNlR4lMvV_AVFgVjZmKMrMWGVG7qhNu",
+        "NERF_FILE_ID": "17dj0pQieo94TozFv-noSBkXebduij1aM",
         "DIET_NERF_MODEL_NAME": "diet_nerf_hotdog",
         "NERF_MODEL_NAME": "nerf_hotdog"
     },
     "Mic": {
-        "FILE_ID": "1vuyCAY8YKLo7OSlTpLJgMem47gJc1W2S",
+        "DIET_NERF_FILE_ID": "1vuyCAY8YKLo7OSlTpLJgMem47gJc1W2S",
+        "NERF_FILE_ID": "17dj0pQieo94TozFv-noSBkXebduij1aM",
         "DIET_NERF_MODEL_NAME": "diet_nerf_mic",
         "NERF_MODEL_NAME": "nerf_mic"
     }