Fix Drums option

app.py

@@ -16,9 +16,10 @@ with open("config.json") as f:
 
 MODEL_DIR = "models"
 
+SCENES_LIST = ["Mic", "Chair", "Lego", "Drums", "Ship", "Hotdog"]
 
-def select_model():
-    obj_select = st.selectbox("Select a Scene", ["Mic", "Chair", "Lego", "Ship", "Hotdog", "Drums"])
+
+def select_model(obj_select):
     DIET_NERF_MODEL_NAME = cfg[obj_select]["DIET_NERF_MODEL_NAME"]
     DIET_NERF_FILE_ID = cfg[obj_select]["DIET_NERF_FILE_ID"]
 
@@ -37,7 +38,8 @@ caption = (
 )
 st.markdown(caption)
 st.markdown("")
-DIET_NERF_MODEL_NAME, DIET_NERF_FILE_ID, NERF_MODEL_NAME, NERF_FILE_ID = select_model()
+obj_select = st.selectbox("Select a Scene", SCENES_LIST, index=0)
+DIET_NERF_MODEL_NAME, DIET_NERF_FILE_ID, NERF_MODEL_NAME, NERF_FILE_ID = select_model(obj_select)
 
 
 @st.cache(show_spinner=False)
@@ -116,7 +118,9 @@ radius = st.sidebar.slider(
 st.markdown("")
 
 with st.spinner("Rendering View..."):
-    with st.spinner("It may take around 1-2 mins. In the meantime, why don't you take a look at our report if you haven't already :)"):
+    with st.spinner(
+        "It may take around 1-2 mins. In the meantime, why don't you take a look at our report 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
@@ -130,14 +134,10 @@ with st.spinner("Rendering View..."):
         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.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.markdown("""<h4 style='text-align: center'>NeRF</h4>""", unsafe_allow_html=True)
         nerf_col.image(n_im, use_column_width=True)
 
         st.markdown(