Using phrases for plotting

app.py

@@ -15,9 +15,9 @@ from segment_anything import SamPredictor
 from supervision.detection.utils import mask_to_polygons
 from supervision.detection.utils import xywh_to_xyxy
 
-if os.environ.get('IS_MY_DEBUG') is None:
-    result = subprocess.run(['pip', 'install', '-e', 'GroundingDINO'], check=True)
-    print(f'pip install GroundingDINO = {result}')
+if os.environ.get("IS_MY_DEBUG") is None:
+    result = subprocess.run(["pip", "install", "-e", "GroundingDINO"], check=True)
+    print(f"pip install GroundingDINO = {result}")
 
 sys.path.append("tag2text")
 sys.path.append("GroundingDINO")
@@ -80,7 +80,9 @@ sam_predictor = SamPredictor(sam)
 sam_automask_generator = SamAutomaticMaskGenerator(sam)
 
 grounding_dino_model = DinoModel(
-    model_config_path=dino_config_file, model_checkpoint_path=dino_checkpoint, device=device
+    model_config_path=dino_config_file,
+    model_checkpoint_path=dino_checkpoint,
+    device=device,
 )
 
 
@@ -104,7 +106,7 @@ def process(image_path, task, prompt, box_threshold, text_threshold, iou_thresho
         metadata["image"]["height"] = h
 
         # Generate tags
-        if task in ["auto", "detection"] and prompt == "":
+        if task in ["auto", "detect"] and prompt == "":
             tags, caption = generate_tags(tag2text_model, image_pil, "None", device)
             prompt = " . ".join(tags)
             print(f"Caption: {caption}")
@@ -129,12 +131,13 @@ def process(image_path, task, prompt, box_threshold, text_threshold, iou_thresho
                 iou_threshold=iou_threshold,
                 post_process=True,
             )
+            print(phrases)
 
             # Draw boxes
             box_annotator = sv.BoxAnnotator()
             labels = [
-                f"{classes[class_id] if class_id else 'Unkown'} {confidence:0.2f}"
-                for _, _, confidence, class_id, _ in detections
+                f"{phrases[i]} {detections.confidence[i]:0.2f}"
+                for i in range(len(phrases))
             ]
             image = box_annotator.annotate(
                 scene=image, detections=detections, labels=labels
@@ -179,24 +182,24 @@ def process(image_path, task, prompt, box_threshold, text_threshold, iou_thresho
 
         # ToDo: Extract metadata
         if detections:
-            id = 1
-            for (xyxy, mask, confidence, class_id, _), area, box_area in zip(
+            i = 0
+            for (xyxy, mask, confidence, _, _), area, box_area in zip(
                 detections, detections.area, detections.box_area
             ):
                 annotation = {
-                    "id": id,
+                    "id": i + 1,
                     "bbox": [int(x) for x in xyxy],
                     "box_area": float(box_area),
                 }
-                if class_id:
-                    annotation["box_confidence"] = float(confidence)
-                    annotation["label"] = classes[class_id] if class_id else "Unkown"
+                if confidence:
+                    annotation["confidence"] = float(confidence)
+                    annotation["label"] = phrases[i]
                 if mask is not None:
                     # annotation["segmentation"] = mask_to_polygons(mask)
                     annotation["area"] = int(area)
-                    annotation["predicted_iou"] = float(scores[id - 1])
+                    annotation["predicted_iou"] = float(scores[i])
                 metadata["annotations"].append(annotation)
-                id += 1
+                i += 1
 
         meta_file = tempfile.NamedTemporaryFile(delete=False, suffix=".json")
         meta_file_path = meta_file.name