add caption + grounding tasks

app.py

@@ -135,6 +135,33 @@ def process_image(image, task_prompt, text_input=None, model_id='microsoft/Flore
         task_prompt = '<MORE_DETAILED_CAPTION>'
         results = run_example(task_prompt, image, model_id=model_id)
         return results, None
+    elif task_prompt == 'Caption + Grounding':
+        task_prompt = '<CAPTION>'
+        results = run_example(task_prompt, image, model_id=model_id)
+        text_input = results[task_prompt]
+        task_prompt = '<CAPTION_TO_PHRASE_GROUNDING>'
+        results = run_example(task_prompt, image, text_input, model_id)
+        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, model_id=model_id)
+        text_input = results[task_prompt]
+        task_prompt = '<CAPTION_TO_PHRASE_GROUNDING>'
+        results = run_example(task_prompt, image, text_input, model_id)
+        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, model_id=model_id)
+        text_input = results[task_prompt]
+        task_prompt = '<CAPTION_TO_PHRASE_GROUNDING>'
+        results = run_example(task_prompt, image, text_input, model_id)
+        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, model_id=model_id)
@@ -210,7 +237,8 @@ with gr.Blocks(css=css) as demo:
                 input_img = gr.Image(label="Input Picture")
                 model_selector = gr.Dropdown(choices=list(models.keys()), label="Model", value='microsoft/Florence-2-large')
                 task_prompt = gr.Dropdown(choices=[
-                    'Caption', 'Detailed Caption', 'More Detailed Caption', 'Object Detection',
+                    'Caption', 'Detailed Caption', 'More Detailed Caption', 'Caption + Grounding',
+                    'Detailed Caption + Grounding', 'More Detailed Caption + Grounding', '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',
@@ -236,4 +264,4 @@ with gr.Blocks(css=css) as demo:
 
         submit_btn.click(process_image, [input_img, task_prompt, text_input, model_selector], [output_text, output_img])
 
-demo.launch(debug=True)
+demo.launch(debug=True)