updated to render multiple objects

app.py

@@ -11,6 +11,22 @@ 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
+
 
 @st.cache
 def download_model():
@@ -34,6 +50,7 @@ model_path = os.path.join(MODEL_DIR, MODEL_NAME)
 if not os.path.isfile(model_path):
     download_model()
 
+MODEL_DIR,MODEL_NAME,FILE_ID = select_model()
 model, state = fetch_model()
 pi = math.pi
 st.sidebar.image("images/diet-nerf.png", width=310)
@@ -55,4 +72,4 @@ 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)
     
-    st.image(im, use_column_width=False)
+    st.image(im, use_column_width=True)