Update app.py

app.py

@@ -1,11 +1,11 @@
+import os
 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
@@ -15,6 +15,16 @@ import matplotlib.patches as patches
 import random
 import numpy as np
 
+from transformers import AutoProcessor, AutoModelForCausalLM
+from detectron2.config import get_cfg
+from detectron2.engine import DefaultPredictor
+from detectron2 import model_zoo
+from detectron2.utils.visualizer import Visualizer, ColorMode
+from detectron2.data import MetadataCatalog
+import cv2
+import numpy as np
+import matplotlib.pyplot as plt
+
 import subprocess
 subprocess.run('pip install flash-attn --no-build-isolation', env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"}, shell=True)
 
@@ -130,11 +140,7 @@ 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')
-example_images = [f for f in os.listdir(example_image_dir) if f.endswith('.jpg')]
-
+# Detectron2 configuration
 cfg = get_cfg()
 cfg.merge_from_file(model_zoo.get_config_file("COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml"))
 cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.5
@@ -142,18 +148,13 @@ cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url("COCO-InstanceSegmentation/mask
 cfg.MODEL.DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 predictor = DefaultPredictor(cfg)
 
-def process_image(image_name):
+def process_image_with_detectron2(image_name):
     image_path = os.path.join(example_image_dir, image_name)
     image = cv2.imread(image_path)
     outputs = predictor(image)
     v = Visualizer(image[:, :, ::-1], MetadataCatalog.get(cfg.DATASETS.TRAIN[0]), scale=1.2)
     out = v.draw_instance_predictions(outputs["instances"].to("cpu"))
-    plt.figure(figsize=(14, 10))
-    plt.imshow(out.get_image()[:, :, ::-1])
-    plt.axis('off')
-    plt.show()
-    return out.get_image()[:, :, ::-1]
-
+    return Image.fromarray(out.get_image()[:, :, ::-1])
 
 with gr.Blocks(theme="sudeepshouche/minimalist") as demo:
     gr.Markdown("## 🧬OmniScience - building teams of fine tuned VLM models for diagnosis and detection 🔧")
@@ -163,6 +164,7 @@ with gr.Blocks(theme="sudeepshouche/minimalist") as demo:
     gr.Markdown("BCCD Datasets on Hugging Face:")
     gr.Markdown("- [🌺 Florence 2](https://huggingface.co/datasets/dwb2023/roboflow100-bccd-florence2/viewer/default/test?q=BloodImage_00038_jpg.rf.1b0ce1635e11b3b49302de527c86bb02.jpg), [💎 PaliGemma](https://huggingface.co/datasets/dwb2023/roboflow-bccd-paligemma/viewer/default/test?q=BloodImage_00038_jpg.rf.1b0ce1635e11b3b49302de527c86bb02.jpg)")
 
+
     with gr.Tab(label="Florence-2 Object Detection"):
         with gr.Row():
             with gr.Column():
@@ -204,12 +206,7 @@ with gr.Blocks(theme="sudeepshouche/minimalist") as demo:
             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)
+        submit_btn.click(process_segmentation, inputs=[input_img], outputs=[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!")
@@ -218,4 +215,4 @@ with gr.Blocks(theme="sudeepshouche/minimalist") as demo:
     gr.Markdown("- [OmniScience fork of Landing AI repo](https://huggingface.co/spaces/dwb2023/omniscience) - I had a lot of fun with this one... some great 🔍reverse engineering enabled by W&B's Weave📊.")
     gr.Markdown("- [Scooby Snacks🐕  - microservice based function calling with style](https://huggingface.co/spaces/dwb2023/blackbird-app) - Leveraging 🤖Claude Sonnet 3.5 to orchestrate Microservice-Based Function Calling.")
 
-demo.launch(debug=True)
+demo.launch(debug=True)