Update app.py

app.py

@@ -2,9 +2,10 @@ import spaces
 import gradio as gr
 from transformers import AutoProcessor, AutoModelForCausalLM
 # import peft
-
+from pycocotools.coco import COCO
 import requests
 import copy
+import os
 
 from PIL import Image, ImageDraw, ImageFont 
 import io
@@ -125,10 +126,31 @@ def process_image(image, task_prompt, text_input=None, model_id='dwb2023/florenc
     else:
         return "", None  # Return empty string and None for unknown task prompts
 
+
+def process_segmentation(image_path, annotation_file_path):
+    coco = COCO(annotation_file_path)
+    img_id = list(coco.imgs.keys())[0]
+    img_info = coco.loadImgs(img_id)[0]
+    ann_ids = coco.getAnnIds(imgIds=img_id)
+    anns = coco.loadAnns(ann_ids)
+
+    image = Image.open(image_path)
+    for ann in anns:
+        if 'segmentation' in ann and ann['segmentation']:
+            image = draw_polygons(image, {'polygons': ann['segmentation'], 'labels': [coco.loadCats(ann['category_id'])[0]['name']]})
+    return image
+
 single_task_list =[
     'Object Detection'
 ]
 
+# Define the path to the example images and annotations file
+example_image_dir = 'examples/bccd-test/'
+annotations_file_path = os.path.join(example_image_dir, '_annotations.coco.json')
+
+# Get the list of example image files
+example_images = [f for f in os.listdir(example_image_dir) if f.endswith('.jpg')]
+
 with gr.Blocks(theme="sudeepshouche/minimalist") as demo:
     gr.Markdown("## 🧬OmniScience - building teams of fine tuned VLM models for diagnosis and detection 🔧")
     gr.Markdown("- 🔬Florence-2 Model Proof of Concept, focusing on Object Detection <OD> tasks.")
@@ -171,6 +193,21 @@ with gr.Blocks(theme="sudeepshouche/minimalist") as demo:
 
         submit_btn.click(process_image, [input_img, task_prompt, model_selector], [output_text, output_img])
 
+    with gr.Tab(label="Segmentation"):
+        with gr.Row():
+            with gr.Column():
+                input_img = gr.Dropdown(choices=example_images, label="Input Picture", value=example_images[0])
+                submit_btn = gr.Button(value="Submit")
+            with gr.Column():
+                output_img = gr.Image(label="Output Image")
+
+        def process_segmentation_tab(image_name):
+            image_path = os.path.join(example_image_dir, image_name)
+            result_image = process_segmentation(image_path, annotations_file_path)
+            return result_image
+
+        submit_btn.click(process_segmentation_tab, input_img, output_img)
+    
     gr.Markdown("## 🚀Other Cool Stuff:")
     gr.Markdown("- [Florence 2 Whitepaper](https://arxiv.org/pdf/2311.06242) - how I found out about the Roboflow 100 and the BCCD dataset.  Turns out this nugget was from the original [Florence whitepaper](https://arxiv.org/pdf/2111.11432) but useful all the same!")
     gr.Markdown("- [Roboflow YouTube Video on Florence 2 fine-tuning](https://youtu.be/i3KjYgxNH6w?si=x1ZMg9hsNe25Y19-&t=1296) - bookmarked an 🧠insightful trade-off analysis of various VLMs.")