full end to end app rewrite

app.py

@@ -1,13 +1,19 @@
+from typing import Tuple, Optional
+
 import gradio as gr
+import spaces
 import supervision as sv
 import torch
-import spaces
+from gradio_image_prompter import ImagePrompter
+from PIL import Image
 
 from utils.annotate import annotate_with_boxes
 from utils.models import load_models, run_inference, CHECKPOINTS
 from utils.tasks import TASK_NAMES, TASKS, OBJECT_DETECTION_TASK_NAME, \
     CAPTION_TASK_NAMES, CAPTION_TASK_NAME, DETAILED_CAPTION_TASK_NAME, \
-    MORE_DETAILED_CAPTION_TASK_NAME, OCR_WITH_REGION_TASK_NAME, OCR_TASK_NAME
+    MORE_DETAILED_CAPTION_TASK_NAME, OCR_WITH_REGION_TASK_NAME, OCR_TASK_NAME, \
+    IMAGE_INPUT_TASK_NAMES, IMAGE_PROMPTER_INPUT_TASK_NAMES, IMAGE_OUTPUT_TASK_NAMES, \
+    TEXTBOX_OUTPUT_TASK_NAMES
 
 MARKDOWN = """
 # Better Florence-2 Playground 🔥
@@ -35,21 +41,14 @@ text format. It uses a DaViT vision encoder to convert images into visual token
 embeddings. These are then concatenated with BERT-generated text embeddings and 
 processed by a transformer-based multi-modal encoder-decoder to generate the response.
 """
-
-OBJECT_DETECTION_EXAMPLES = [
-    ["microsoft/Florence-2-large-ft", OBJECT_DETECTION_TASK_NAME, "https://media.roboflow.com/notebooks/examples/dog-2.jpeg"]
-]
-CAPTION_EXAMPLES = [
-    ["microsoft/Florence-2-large-ft", CAPTION_TASK_NAME, "https://media.roboflow.com/notebooks/examples/dog-2.jpeg"],
-    ["microsoft/Florence-2-large-ft", DETAILED_CAPTION_TASK_NAME, "https://media.roboflow.com/notebooks/examples/dog-2.jpeg"],
-    ["microsoft/Florence-2-large-ft", MORE_DETAILED_CAPTION_TASK_NAME, "https://media.roboflow.com/notebooks/examples/dog-2.jpeg"]
-]
-OCR_EXAMPLES = [
-    ["microsoft/Florence-2-large-ft", OCR_TASK_NAME, "https://media.roboflow.com/notebooks/examples/handwritten-text.jpg"],
-]
-OCR_WITH_REGION_EXAMPLES = [
-    ["microsoft/Florence-2-large-ft", OCR_WITH_REGION_TASK_NAME, "https://media.roboflow.com/notebooks/examples/handwritten-text.jpg"],
-    ["microsoft/Florence-2-large-ft", OCR_WITH_REGION_TASK_NAME, "https://media.roboflow.com/inference/license_plate_1.jpg"]
+EXAMPLES = [
+    ["microsoft/Florence-2-large-ft", OBJECT_DETECTION_TASK_NAME, "https://media.roboflow.com/notebooks/examples/dog-2.jpeg", None],
+    ["microsoft/Florence-2-large-ft", CAPTION_TASK_NAME, "https://media.roboflow.com/notebooks/examples/dog-2.jpeg", None],
+    ["microsoft/Florence-2-large-ft", DETAILED_CAPTION_TASK_NAME, "https://media.roboflow.com/notebooks/examples/dog-2.jpeg", None],
+    ["microsoft/Florence-2-large-ft", MORE_DETAILED_CAPTION_TASK_NAME, "https://media.roboflow.com/notebooks/examples/dog-2.jpeg", None],
+    ["microsoft/Florence-2-large-ft", OCR_TASK_NAME, "https://media.roboflow.com/notebooks/examples/handwritten-text.jpg", None],
+    ["microsoft/Florence-2-large-ft", OCR_WITH_REGION_TASK_NAME, "https://media.roboflow.com/notebooks/examples/handwritten-text.jpg", None],
+    ["microsoft/Florence-2-large-ft", OCR_WITH_REGION_TASK_NAME, "https://media.roboflow.com/inference/license_plate_1.jpg", None]
 ]
 
 DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
@@ -58,20 +57,26 @@ MODELS, PROCESSORS = load_models(DEVICE)
 
 
 @spaces.GPU
-def process(checkpoint_dropdown, task_dropdown, image_input):
+def process(
+    checkpoint_dropdown,
+    task_dropdown,
+    image_input,
+    image_prompter_input
+) -> Tuple[Optional[Image.Image], Optional[str]]:
     model = MODELS[checkpoint_dropdown]
     processor = PROCESSORS[checkpoint_dropdown]
     task = TASKS[task_dropdown]
+
     if task_dropdown in [OBJECT_DETECTION_TASK_NAME, OCR_WITH_REGION_TASK_NAME]:
         _, response = run_inference(
             model, processor, DEVICE, image_input, task)
         detections = sv.Detections.from_lmm(
             lmm=sv.LMM.FLORENCE_2, result=response, resolution_wh=image_input.size)
-        return annotate_with_boxes(image_input, detections)
+        return annotate_with_boxes(image_input, detections), None
     elif task_dropdown in CAPTION_TASK_NAMES or task_dropdown == OCR_TASK_NAME:
         _, response = run_inference(
             model, processor, DEVICE, image_input, task)
-        return response[task]
+        return None, response[task]
 
 
 with gr.Blocks() as demo:
@@ -80,31 +85,74 @@ with gr.Blocks() as demo:
         checkpoint_dropdown_component = gr.Dropdown(
             choices=CHECKPOINTS,
             value=CHECKPOINTS[0],
-            label="Model", info="Select a Florence 2 model to use.")
+            label="Model", info="Select a Florence 2 model to use.",
+            interactive=True
+        )
         task_dropdown_component = gr.Dropdown(
             choices=TASK_NAMES,
             value=TASK_NAMES[0],
-            label="Task", info="Select a task to perform with the model.")
+            label="Task", info="Select a task to perform with the model.",
+            interactive=True
+        )
 
     with gr.Row():
         with gr.Column():
-            image_input_component = gr.Image(type='pil', label='Image Input')
+            image_input_component = gr.Image(
+                type='pil', label='Upload image')
+            image_prompter_input_component = ImagePrompter(
+                type='pil', label='Upload image and draw box prompt', visible=False)
             submit_button_component = gr.Button(value='Submit', variant='primary')
 
         with gr.Column():
             image_output_component = gr.Image(type='pil', label='Image Output')
             text_output_component = gr.Textbox(label='Caption Output', visible=False)
-
-    def on_dropdown_input(text):
-        if text in CAPTION_TASK_NAMES + [OCR_TASK_NAME]:
-            return [gr.Image(visible=False), gr.Textbox(visible=True)]
-        else:
-            return [gr.Image(visible=True), gr.Textbox(visible=False)]
-
-    task_dropdown_component.input(
-        on_dropdown_input,
+    with gr.Row():
+        gr.Examples(
+            fn=process,
+            examples=EXAMPLES,
+            inputs=[
+                checkpoint_dropdown_component,
+                task_dropdown_component,
+                image_input_component,
+                image_prompter_input_component
+            ],
+            outputs=[
+                image_output_component,
+                text_output_component
+            ],
+            run_on_click=True
+        )
+
+    def on_dropdown_change(text):
+        return [
+            gr.Image(visible=text in IMAGE_INPUT_TASK_NAMES),
+            ImagePrompter(visible=text in IMAGE_PROMPTER_INPUT_TASK_NAMES),
+            gr.Image(visible=text in IMAGE_OUTPUT_TASK_NAMES),
+            gr.Textbox(visible=text in TEXTBOX_OUTPUT_TASK_NAMES)
+        ]
+
+    task_dropdown_component.change(
+        on_dropdown_change,
         inputs=[task_dropdown_component],
-        outputs=[image_output_component, text_output_component])
-
+        outputs=[
+            image_input_component,
+            image_prompter_input_component,
+            image_output_component,
+            text_output_component
+        ]
+    )
+    submit_button_component.click(
+        fn=process,
+        inputs=[
+            checkpoint_dropdown_component,
+            task_dropdown_component,
+            image_input_component,
+            image_prompter_input_component
+        ],
+        outputs=[
+            image_output_component,
+            text_output_component
+        ]
+    )
 
 demo.launch(debug=False, show_error=True, max_threads=1)

utils/tasks.py

@@ -4,6 +4,8 @@ DETAILED_CAPTION_TASK_NAME = "Detailed Caption"
 MORE_DETAILED_CAPTION_TASK_NAME = "More Detailed Caption"
 OCR_TASK_NAME = "OCR"
 OCR_WITH_REGION_TASK_NAME = "OCR with Region"
+REGION_TO_CATEGORY_TASK_NAME = "Region to Category"
+REGION_TO_DESCRIPTION_TASK_NAME = "Region to Description"
 
 TASK_NAMES = [
     OBJECT_DETECTION_TASK_NAME,
@@ -11,7 +13,9 @@ TASK_NAMES = [
     DETAILED_CAPTION_TASK_NAME,
     MORE_DETAILED_CAPTION_TASK_NAME,
     OCR_TASK_NAME,
-    OCR_WITH_REGION_TASK_NAME
+    OCR_WITH_REGION_TASK_NAME,
+    REGION_TO_CATEGORY_TASK_NAME,
+    REGION_TO_DESCRIPTION_TASK_NAME
 ]
 TASKS = {
     OBJECT_DETECTION_TASK_NAME: "<OD>",
@@ -19,10 +23,37 @@ TASKS = {
     DETAILED_CAPTION_TASK_NAME: "<DETAILED_CAPTION>",
     MORE_DETAILED_CAPTION_TASK_NAME: "<MORE_DETAILED_CAPTION>",
     OCR_TASK_NAME: "<OCR>",
-    OCR_WITH_REGION_TASK_NAME: "<OCR_WITH_REGION>"
+    OCR_WITH_REGION_TASK_NAME: "<OCR_WITH_REGION>",
+    REGION_TO_CATEGORY_TASK_NAME: "<REGION_TO_CATEGORY>",
+    REGION_TO_DESCRIPTION_TASK_NAME: "<REGION_TO_DESCRIPTION>"
 }
 CAPTION_TASK_NAMES = [
     CAPTION_TASK_NAME,
     DETAILED_CAPTION_TASK_NAME,
     MORE_DETAILED_CAPTION_TASK_NAME
+]
+
+IMAGE_INPUT_TASK_NAMES = [
+    OBJECT_DETECTION_TASK_NAME,
+    CAPTION_TASK_NAME,
+    DETAILED_CAPTION_TASK_NAME,
+    MORE_DETAILED_CAPTION_TASK_NAME,
+    OCR_TASK_NAME,
+    OCR_WITH_REGION_TASK_NAME,
+]
+IMAGE_PROMPTER_INPUT_TASK_NAMES = [
+    REGION_TO_CATEGORY_TASK_NAME,
+    REGION_TO_DESCRIPTION_TASK_NAME
+]
+IMAGE_OUTPUT_TASK_NAMES = [
+    OBJECT_DETECTION_TASK_NAME,
+    OCR_WITH_REGION_TASK_NAME,
+    REGION_TO_CATEGORY_TASK_NAME,
+    REGION_TO_DESCRIPTION_TASK_NAME
+]
+TEXTBOX_OUTPUT_TASK_NAMES = [
+    CAPTION_TASK_NAME,
+    DETAILED_CAPTION_TASK_NAME,
+    MORE_DETAILED_CAPTION_TASK_NAME,
+    OCR_TASK_NAME
 ]