Update app.py

app.py

@@ -34,24 +34,30 @@ def generate_caption(image, model_name):
 
 # Predefined images for selection
 image_folder = "images"
-predefined_images_paths = [
-    os.path.join(image_folder, fname) for fname in os.listdir(image_folder) if fname.lower().endswith(('.png', '.jpg', '.jpeg'))
+predefined_images = [
+    Image.open(os.path.join(image_folder, fname)).convert("RGB")
+    for fname in os.listdir(image_folder) \
+    if  fname.lower().endswith(('.png', '.jpg', '.jpeg', '.gif', '.bmp', '.ppm')) 
 ]
 
-# Gradio app
-def app(image=None, model_name="laicsiifes/swin-distilbertimbau"):
+# Define components logic
+def handle_uploaded_image(image):
     if image is None:
-        return "Please upload a valid image."
-    return generate_caption(image, model_name)
+        return None
+    pil_image = image.convert("RGB")
+    return pil_image, None
+
+def handle_generate_button(image, selected_model):
+    if image is None:
+        return "Please upload an image to generate a caption."
+    return generate_caption(image, selected_model)
 
 # Define UI
 with gr.Blocks(theme=gr.themes.Citrus(primary_hue="blue", secondary_hue="orange")) as interface:
     gr.Markdown("""
         # Welcome to the LAICSI-IFES space for Vision Encoder-Decoder (VED) demonstration
-        
         ---
-        
-        ### Be patient with the Swin-GPorTuguese-2 as it is heavier than the Swin-DistilBERTimbau.
+        ### Select an available model: Swin-DistilBERTimbau (168M) or Swin-GPorTuguese-2 (240M)
     """)
     with gr.Row(variant='panel'):
         with gr.Column():
@@ -61,12 +67,14 @@ with gr.Blocks(theme=gr.themes.Citrus(primary_hue="blue", secondary_hue="orange"
                 label="Select Model"
             )
             
-    gr.Markdown("""---""")
+    gr.Markdown("""
+        ---
+        ### Upload image or example images below, and click `Generate`
+        """)
     
     with gr.Row(variant='panel'):
         with gr.Column():
-            image_display = gr.Image(type="pil", label="Image Preview", interactive=False, height=400)
-            upload_button = gr.File(label="Upload an Image", file_types=["image"], type="filepath", height=100)
+            image_display = gr.Image(type="pil", label="Image Preview", image_mode="RGB", height=400)
         with gr.Column():
             output_text = gr.Textbox(label="Generated Caption")
             generate_button = gr.Button("Generate")
@@ -74,22 +82,20 @@ with gr.Blocks(theme=gr.themes.Citrus(primary_hue="blue", secondary_hue="orange"
     gr.Markdown("""---""")
     
     with gr.Row(variant='panel'):
-            examples = gr.Examples(predefined_images_paths, inputs=[upload_button], label="Examples")
+            examples = gr.Examples(
+                examples=predefined_images,
+                fn=handle_uploaded_image,
+                inputs=[image_display],
+                outputs=[image_display, output_text],
+                label="Examples"
+            )
 
-    # Define logic
-    def handle_uploaded_image(image):
-        if image is None:
-            return None
-        pil_image = Image.open(image).convert("RGB")
-        return pil_image
+    # Define actions
+    model_selector.change(fn=lambda: (None, None), outputs=[image_display, output_text])
 
-    def handle_generate_button(image, selected_model):
-        if image is None:
-            return "Please upload an image to generate a caption."
-        return generate_caption(image, selected_model)
+    image_display.upload(fn=handle_uploaded_image, inputs=[image_display], outputs=[image_display, output_text])
+    image_display.clear(fn=lambda: None, outputs=[output_text])
 
-    model_selector.change(fn=lambda _: (None, None, None), inputs=[model_selector], outputs=[image_display, upload_button, output_text])
-    upload_button.change(fn=handle_uploaded_image, inputs=upload_button, outputs=image_display)
     generate_button.click(fn=handle_generate_button, inputs=[image_display, model_selector], outputs=output_text)
 
-    interface.launch(share=False)
+    interface.launch(share=False)