Added gokaygokay's "Cascaded task"

app.py

@@ -1,3 +1,6 @@
+import os
+os.environ["PYTORCH_ENABLE_MPS_FALLBACK"] = "1"
+
 import gradio as gr
 from transformers import AutoProcessor, AutoModelForCausalLM
 import spaces
@@ -13,7 +16,7 @@ import matplotlib.patches as patches
 import random
 import numpy as np
 
-import os
+
 from unittest.mock import patch
 from transformers import AutoModelForCausalLM, AutoProcessor
 from transformers.dynamic_module_utils import get_imports
@@ -145,16 +148,43 @@ def process_image(image, task_prompt, text_input=None):
     image = Image.fromarray(image)  # Convert NumPy array to PIL Image
     if task_prompt == 'Caption':
         task_prompt = '<CAPTION>'
-        result = run_example(task_prompt, image)
-        return result, None
+        results = run_example(task_prompt, image)
+        return results, None
     elif task_prompt == 'Detailed Caption':
         task_prompt = '<DETAILED_CAPTION>'
-        result = run_example(task_prompt, image)
-        return result, None
+        results = run_example(task_prompt, image)
+        return results, None
     elif task_prompt == 'More Detailed Caption':
         task_prompt = '<MORE_DETAILED_CAPTION>'
-        result = run_example(task_prompt, image)
-        return result, None
+        results = run_example(task_prompt, image)
+        return results, None
+    elif task_prompt == 'Caption + Grounding':
+        task_prompt = '<CAPTION>'
+        results = run_example(task_prompt, image)
+        text_input = results[task_prompt]
+        task_prompt = '<CAPTION_TO_PHRASE_GROUNDING>'
+        results = run_example(task_prompt, image, text_input)
+        results['<CAPTION>'] = text_input
+        fig = plot_bbox(image, results['<CAPTION_TO_PHRASE_GROUNDING>'])
+        return results, fig_to_pil(fig)
+    elif task_prompt == 'Detailed Caption + Grounding':
+        task_prompt = '<DETAILED_CAPTION>'
+        results = run_example(task_prompt, image)
+        text_input = results[task_prompt]
+        task_prompt = '<CAPTION_TO_PHRASE_GROUNDING>'
+        results = run_example(task_prompt, image, text_input)
+        results['<DETAILED_CAPTION>'] = text_input
+        fig = plot_bbox(image, results['<CAPTION_TO_PHRASE_GROUNDING>'])
+        return results, fig_to_pil(fig)
+    elif task_prompt == 'More Detailed Caption + Grounding':
+        task_prompt = '<MORE_DETAILED_CAPTION>'
+        results = run_example(task_prompt, image)
+        text_input = results[task_prompt]
+        task_prompt = '<CAPTION_TO_PHRASE_GROUNDING>'
+        results = run_example(task_prompt, image, text_input)
+        results['<MORE_DETAILED_CAPTION>'] = text_input
+        fig = plot_bbox(image, results['<CAPTION_TO_PHRASE_GROUNDING>'])
+        return results, fig_to_pil(fig)
     elif task_prompt == 'Object Detection':
         task_prompt = '<OD>'
         results = run_example(task_prompt, image)
@@ -203,8 +233,8 @@ def process_image(image, task_prompt, text_input=None):
         return results, None
     elif task_prompt == 'OCR':
         task_prompt = '<OCR>'
-        result = run_example(task_prompt, image)
-        return result, None
+        results = run_example(task_prompt, image)
+        return results, None
     elif task_prompt == 'OCR with Region':
         task_prompt = '<OCR_WITH_REGION>'
         results = run_example(task_prompt, image)
@@ -222,19 +252,37 @@ css = """
   }
 """
 
+
+single_task_list =[
+    'Caption', 'Detailed Caption', 'More Detailed Caption', 'Object Detection',
+    'Dense Region Caption', 'Region Proposal', 'Caption to Phrase Grounding',
+    'Referring Expression Segmentation', 'Region to Segmentation',
+    'Open Vocabulary Detection', 'Region to Category', 'Region to Description',
+    'OCR', 'OCR with Region'
+]
+
+cascaded_task_list =[
+    'Caption + Grounding', 'Detailed Caption + Grounding', 'More Detailed Caption + Grounding'
+]
+
+
+def update_task_dropdown(choice):
+    if choice == 'Cascaded task':
+        return gr.Dropdown(choices=cascaded_task_list, value='Caption + Grounding')
+    else:
+        return gr.Dropdown(choices=single_task_list, value='Caption')
+
+
+
 with gr.Blocks(css=css) as demo:
     gr.Markdown(DESCRIPTION)
     with gr.Tab(label="Florence-2 Image Captioning"):
         with gr.Row():
             with gr.Column():
-                input_img = gr.Image(label="Input Picture")
-                task_prompt = gr.Dropdown(choices=[
-                    'Caption', 'Detailed Caption', 'More Detailed Caption', 'Object Detection',
-                    'Dense Region Caption', 'Region Proposal', 'Caption to Phrase Grounding',
-                    'Referring Expression Segmentation', 'Region to Segmentation',
-                    'Open Vocabulary Detection', 'Region to Category', 'Region to Description',
-                    'OCR', 'OCR with Region'
-                ], label="Task Prompt", value= 'Caption')
+                input_img = gr.Image(label="Input Picture")                
+                task_type = gr.Radio(choices=['Single task', 'Cascaded task'], label='Task type selector', value='Single task')
+                task_prompt = gr.Dropdown(choices=single_task_list, label="Task Prompt", value="Caption")
+                task_type.change(fn=update_task_dropdown, inputs=task_type, outputs=task_prompt)
                 text_input = gr.Textbox(label="Text Input (optional)")
                 submit_btn = gr.Button(value="Submit")
             with gr.Column():
@@ -255,4 +303,4 @@ with gr.Blocks(css=css) as demo:
 
         submit_btn.click(process_image, [input_img, task_prompt, text_input], [output_text, output_img])
 
-demo.launch(debug=True)
+demo.launch(debug=True)