small updates

.gitattributes

@@ -32,3 +32,5 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+*.jpg filter=lfs diff=lfs merge=lfs -text
+*.png filter=lfs diff=lfs merge=lfs -text

app.py

@@ -4,6 +4,15 @@ import os
 from PIL import Image
 from functools import partial
 
+def retrieve_input_image_wild(dataset, inputs):
+    img_id = inputs
+    img_path = os.path.join('online_demo', dataset, 'step-100_scale-6.0')
+    try:
+        image = Image.open(os.path.join(img_path, '%s.jpg' % img_id))
+    except:
+        image = Image.open(os.path.join(img_path, '%s.png' % img_id))
+    return image
+
 def retrieve_input_image(dataset, inputs):
     img_id = inputs
     img_path = os.path.join('online_demo', dataset, 'step-100_scale-6.0', img_id, 'input.png')
@@ -25,12 +34,12 @@ with gr.Blocks() as demo:
     with gr.Tab("In-the-wild Images"):
         with gr.Row():
             with gr.Column(scale=1):
-                default_input_image = Image.open( os.path.join('online_demo', 'nerf_wild', 'step-100_scale-6.0', 'car1', 'input.png'))
+                default_input_image = Image.open( os.path.join('online_demo', 'nerf_wild', 'step-100_scale-6.0', 'car1.png'))
                 input_image = gr.Image(default_input_image, shape=[256, 256])
-                options = sorted(os.listdir('online_demo/nerf_wild/step-100_scale-6.0'))
+                options = sorted(next(os.walk('online_demo/nerf_wild/step-100_scale-6.0'))[1])
                 img_id = gr.Dropdown(options, value='car1', label='options')
                 text_button = gr.Button("Load Input Image")
-                retrieve_input_image_dataset = partial(retrieve_input_image, 'nerf_wild')
+                retrieve_input_image_dataset = partial(retrieve_input_image_wild, 'nerf_wild')
                 text_button.click(retrieve_input_image_dataset, inputs=img_id, outputs=input_image)
 
             with gr.Column(scale=1):
@@ -41,7 +50,7 @@ with gr.Blocks() as demo:
                 gr.Slider(-0.5, 0.5, value=0, step=0.5, label='Zoom'),
                 gr.Slider(0, 3, value=1, step=1, label='Random seed')]
                 
-                submit_button = gr.Button("Get Novel View")
+                submit_button = gr.Button("Generate Novel View")
                 retrieve_novel_view_dataset = partial(retrieve_novel_view, 'nerf_wild')
                 submit_button.click(retrieve_novel_view_dataset, inputs=inputs, outputs=novel_view)
         
@@ -64,7 +73,7 @@ with gr.Blocks() as demo:
                 gr.Slider(-0.5, 0.5, value=0, step=0.5, label='Zoom'),
                 gr.Slider(0, 3, value=1, step=1, label='Random seed')]
                 
-                submit_button = gr.Button("Get Novel View")
+                submit_button = gr.Button("Generate Novel View")
                 retrieve_novel_view_dataset = partial(retrieve_novel_view, 'GSO')
                 submit_button.click(retrieve_novel_view_dataset, inputs=inputs, outputs=novel_view)