Comparison of DietNeRF and NeRF

app.py

@@ -21,7 +21,11 @@ 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"]
-    return DIET_NERF_MODEL_NAME, 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")
@@ -33,29 +37,48 @@ caption = (
 )
 st.markdown(caption)
 st.markdown("")
-DIET_NERF_MODEL_NAME, DIET_NERF_FILE_ID = select_model()
+DIET_NERF_MODEL_NAME, DIET_NERF_FILE_ID, NERF_MODEL_NAME, NERF_FILE_ID = select_model()
 
 
 @st.cache(show_spinner=False)
-def download_model():
+def download_diet_nerf_model():
     os.makedirs(MODEL_DIR, exist_ok=True)
-    _model_path = os.path.join(MODEL_DIR, DIET_NERF_MODEL_NAME)
+    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, _model_path, quiet=False)
-    print(f"Model downloaded from google drive: {_model_path}")
+    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_model():
+def fetch_diet_nerf_model():
     model, state = load_trained_model(MODEL_DIR, DIET_NERF_MODEL_NAME)
     return model, state
 
 
-model_path = os.path.join(MODEL_DIR, DIET_NERF_MODEL_NAME)
-if not os.path.isfile(model_path):
-    download_model()
+@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()
 
-model, state = fetch_model()
 pi = math.pi
 
 st.sidebar.markdown(
@@ -94,14 +117,25 @@ st.markdown("")
 
 with st.spinner("Rendering View..."):
     with st.spinner("It may take 2-3 mins. So, why don't you read our report in the meantime"):
-        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))
-
-        # diet_nerf_col = st.beta_columns(1)
-        # diet_nerf_col.markdown(
-        #     """<h4 style='text-align: center'>DietNerF</h4>""", unsafe_allow_html=True
-        # )
-        st.image(im, use_column_width="auto")
+        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)

config.json

@@ -7,25 +7,25 @@
     },
     "Lego": {
         "DIET_NERF_FILE_ID": "1UaUKgShJTKp5zNQdXvVhkyxgU-2YU3Ay",
-        "NERF_FILE_ID": "17dj0pQieo94TozFv-noSBkXebduij1aM",
+        "NERF_FILE_ID": "16_Mx_p7Sncbb5MoP7JypXAhvmfNfnOwf",
         "DIET_NERF_MODEL_NAME": "diet_nerf_lego",
         "NERF_MODEL_NAME": "nerf_lego"
     },
     "Ship": {
         "DIET_NERF_FILE_ID": "1Eb4eZ9vn6FcbqLIaBoW-6NlyWr-y0eTW",
-        "NERF_FILE_ID": "17dj0pQieo94TozFv-noSBkXebduij1aM",
+        "NERF_FILE_ID": "1fkChM8-44uSsMcMu6XIFB33nXf58A-Il",
         "DIET_NERF_MODEL_NAME": "diet_nerf_ship",
         "NERF_MODEL_NAME": "nerf_ship"
     },
     "Hotdog": {
         "DIET_NERF_FILE_ID": "11vNlR4lMvV_AVFgVjZmKMrMWGVG7qhNu",
-        "NERF_FILE_ID": "17dj0pQieo94TozFv-noSBkXebduij1aM",
+        "NERF_FILE_ID": "11vNlR4lMvV_AVFgVjZmKMrMWGVG7qhNu",
         "DIET_NERF_MODEL_NAME": "diet_nerf_hotdog",
         "NERF_MODEL_NAME": "nerf_hotdog"
     },
     "Mic": {
         "DIET_NERF_FILE_ID": "1vuyCAY8YKLo7OSlTpLJgMem47gJc1W2S",
-        "NERF_FILE_ID": "17dj0pQieo94TozFv-noSBkXebduij1aM",
+        "NERF_FILE_ID": "1JK3vu2F9xnOuhPjYVflMkPp0BcAc4w_9",
         "DIET_NERF_MODEL_NAME": "diet_nerf_mic",
         "NERF_MODEL_NAME": "nerf_mic"
     }