Update app.py

app.py

@@ -126,20 +126,6 @@ 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'
 ]
@@ -147,10 +133,28 @@ single_task_list =[
 # 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')]
 
+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
+cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url("COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml")
+cfg.MODEL.DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+predictor = DefaultPredictor(cfg)
+
+def process_image(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]
+
+
 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.")
@@ -159,7 +163,6 @@ 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():