add interface logic

app.py

@@ -1,15 +1,13 @@
 import gradio as gr
 import torch
-from PIL import Image, ImageDraw, ImageFont 
+from PIL import Image, ImageDraw
 from transformers import AutoProcessor
 from modeling_florence2 import Florence2ForConditionalGeneration
-from configuration_florence2 import Florence2Config
 import io
 import matplotlib.pyplot as plt
 import matplotlib.patches as patches
 import numpy as np
 import random
-import copy
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
 torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32
@@ -28,6 +26,9 @@ TASK_PROMPTS = {
     "Region Proposal": "<REGION_PROPOSAL>"
 }
 
+IMAGE_TASKS = ["Object Detection", "Dense Region Caption", "Region Proposal", "OCR with Region"]
+TEXT_TASKS = ["Caption", "Detailed Caption", "More Detailed Caption", "OCR"]
+
 colormap = ['blue','orange','green','purple','brown','pink','gray','olive','cyan','red',
             'lime','indigo','violet','aqua','magenta','coral','gold','tan','skyblue']
 
@@ -78,11 +79,10 @@ def process_image(image, task):
     
     return parsed_answer
 
-
 def main_process(image, task):
     result = process_image(image, task)
     
-    if task in ["Object Detection", "Dense Region Caption", "Region Proposal", "OCR with Region"]:
+    if task in IMAGE_TASKS:
         if task == "OCR with Region":
             output_image = draw_ocr_bboxes(image.copy(), result[TASK_PROMPTS[task]])
         else:
@@ -93,7 +93,7 @@ def main_process(image, task):
         return None, gr.update(visible=False), str(result), gr.update(visible=True)
 
 def reset_outputs():
-    return None, gr.update(visible=False), None, gr.update(visible=False)
+    return None, gr.update(visible=False), None, gr.update(visible=True)
 
 with gr.Blocks(title="Florence-2 Demo") as iface:
     gr.Markdown("# Florence-2 Demo")
@@ -101,17 +101,22 @@ with gr.Blocks(title="Florence-2 Demo") as iface:
     
     with gr.Column():
         image_input = gr.Image(type="pil", label="Input Image")
-        task_dropdown = gr.Dropdown(list(TASK_PROMPTS.keys()), label="Task")
+        task_dropdown = gr.Dropdown(list(TASK_PROMPTS.keys()), label="Task", value="Caption")
     
     with gr.Row():
         submit_button = gr.Button("Process")
         reset_button = gr.Button("Reset")
     
     output_image = gr.Image(label="Processed Image", visible=False)
-    output_text = gr.Textbox(label="Output", visible=False)
+    output_text = gr.Textbox(label="Output", visible=True)
+    
+    def process_and_update(image, task):
+        if image is None:
+            return None, gr.update(visible=False), "Please upload an image first.", gr.update(visible=True)
+        return main_process(image, task)
     
     submit_button.click(
-        fn=main_process,
+        fn=process_and_update,
         inputs=[image_input, task_dropdown],
         outputs=[output_image, output_image, output_text, output_text]
     )
@@ -121,5 +126,11 @@ with gr.Blocks(title="Florence-2 Demo") as iface:
         inputs=[],
         outputs=[output_image, output_image, output_text, output_text]
     )
+    
+    task_dropdown.change(
+        fn=lambda task: (gr.update(visible=task in IMAGE_TASKS), gr.update(visible=task in TEXT_TASKS)),
+        inputs=[task_dropdown],
+        outputs=[output_image, output_text]
+    )
 
 iface.launch()