Merge pull request #529 from microsoft/PreRelease

Supervision 0.19 compatibility and #522

PW_FT_classification/src/utils/utils.py

@@ -35,7 +35,7 @@ def save_crop_images(results, output_dir, original_csv_path, overwrite=False):
         for entry in results:
             # Process the data if the name of the file is in the dataframe
             if os.path.basename(entry["img_id"]) in original_df['path'].values:
-                for i, (xyxy, _, _, cat, _) in enumerate(entry["detections"]):
+                for i, (xyxy, cat) in enumerate(zip(entry["detections"].xyxy, entry["detections"].class_id)):
                     cropped_img = sv.crop_image(
                         image=np.array(Image.open(entry["img_id"]).convert("RGB")), xyxy=xyxy
                     )

PytorchWildlife/models/detection/yolov5/base_detector.py

@@ -86,7 +86,7 @@ def results_generation(self, preds, img_id, id_strip=None):
         )
         results["labels"] = [
             f"{self.CLASS_NAMES[class_id]} {confidence:0.2f}"
-            for _, _, confidence, class_id, _ in results["detections"]
+            for confidence, class_id in zip(results["detections"].confidence, results["detections"].class_id)
         ]
         return results
 

PytorchWildlife/utils/post_process.py

@@ -36,34 +36,33 @@ def save_detection_images(results, output_dir, input_dir = None, overwrite=False
         overwrite (bool):
             Whether overwriting existing image folders. Default to False.
     """
-    box_annotator = sv.BoxAnnotator(thickness=4, text_thickness=4, text_scale=2)
+    box_annotator = sv.BoundingBoxAnnotator(thickness=4)
+    lab_annotator = sv.LabelAnnotator(text_color=sv.Color.BLACK, text_thickness=4, text_scale=2)
     os.makedirs(output_dir, exist_ok=True)
 
-    if isinstance(results, list):
-        for entry in results:
-            annotated_img = box_annotator.annotate(
-                scene=np.array(Image.open(entry["img_id"]).convert("RGB")),
-                detections=entry["detections"],
-                labels=entry["labels"],
-            )
-
-            img_id_parts=Path(entry["img_id"]).parts
-            last_input_dir=Path(input_dir).parts[-1]
-            relative_dir=Path(*img_id_parts[img_id_parts.index(last_input_dir)+1:-1])
-            full_output_dir = os.path.join(output_dir, relative_dir)
-            os.makedirs(full_output_dir, exist_ok=True)
-            with sv.ImageSink(target_dir_path=full_output_dir, overwrite=overwrite) as sink: 
+    with sv.ImageSink(target_dir_path=output_dir, overwrite=True) as sink:
+        if isinstance(results, list):
+            for entry in results:
+                annotated_img = lab_annotator.annotate(
+                    scene=box_annotator.annotate(
+                        scene=np.array(Image.open(entry["img_id"]).convert("RGB")),
+                        detections=entry["detections"],
+                    ),
+                    detections=entry["detections"],
+                    labels=entry["labels"],
+                )
                 sink.save_image(
                     image=cv2.cvtColor(annotated_img, cv2.COLOR_RGB2BGR), image_name=entry["img_id"].rsplit(os.sep, 1)[1]
                 )
-    else:
-        annotated_img = box_annotator.annotate(
-            scene=np.array(Image.open(results["img_id"]).convert("RGB")),
-            detections=results["detections"],
-            labels=results["labels"],
-        )
-
-        with sv.ImageSink(target_dir_path=output_dir, overwrite=overwrite) as sink: 
+        else:
+            annotated_img = lab_annotator.annotate(
+                scene=box_annotator.annotate(
+                    scene=np.array(Image.open(results["img_id"]).convert("RGB")),
+                    detections=results["detections"],
+                ),
+                detections=results["detections"],
+                labels=results["labels"],
+            )
             sink.save_image(
                 image=cv2.cvtColor(annotated_img, cv2.COLOR_RGB2BGR), image_name=results["img_id"].rsplit(os.sep, 1)[1]
             )
@@ -82,29 +81,24 @@ def save_crop_images(results, output_dir, input_dir = None, overwrite=False):
         overwrite (bool):
             Whether overwriting existing image folders. Default to False.
     """
+    if isinstance(results, dict):
+        results = [results]
+
     assert isinstance(results, list)
     os.makedirs(output_dir, exist_ok=True)
-
-    for entry in results:
-        for i, (xyxy, _, _, cat, _) in enumerate(entry["detections"]):
-            cropped_img = sv.crop_image(
-                image=np.array(Image.open(entry["img_id"]).convert("RGB")), xyxy=xyxy
-            )
-
-            img_id_parts=Path(entry["img_id"]).parts
-            last_input_dir=Path(input_dir).parts[-1]
-            relative_dir=Path(*img_id_parts[img_id_parts.index(last_input_dir)+1:-1])
-            full_output_dir = os.path.join(output_dir, relative_dir)
-            os.makedirs(full_output_dir, exist_ok=True)
-            with sv.ImageSink(target_dir_path=full_output_dir, overwrite=overwrite) as sink:
+    with sv.ImageSink(target_dir_path=output_dir, overwrite=True) as sink:
+        for entry in results:
+            for i, (xyxy, cat) in enumerate(zip(entry["detections"].xyxy, entry["detections"].class_id)):
+                cropped_img = sv.crop_image(
+                    image=np.array(Image.open(entry["img_id"]).convert("RGB")), xyxy=xyxy
+                )
                 sink.save_image(
                     image=cv2.cvtColor(cropped_img, cv2.COLOR_RGB2BGR),
                     image_name="{}_{}_{}".format(
                         int(cat), i, entry["img_id"].rsplit(os.sep, 1)[1]
                     ),
                 )
 
-
 def save_detection_json(det_results, output_dir, categories=None, exclude_category_ids=[], exclude_file_path=None):
     """
     Save detection results to a JSON file.

demo/gradio_demo.py

@@ -29,7 +29,8 @@
 # Setting the device to use for computations ('cuda' indicates GPU)
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 # Initializing a supervision box annotator for visualizing detections
-box_annotator = sv.BoxAnnotator(thickness=4, text_thickness=4, text_scale=2)
+box_annotator = sv.BoundingBoxAnnotator(thickness=4)
+lab_annotator = sv.LabelAnnotator(text_color=sv.Color.BLACK, text_thickness=4, text_scale=2)
 # Create a temp folder
 os.makedirs(os.path.join("..","temp"), exist_ok=True)
 
@@ -90,7 +91,14 @@ def single_image_detection(input_img, det_conf_thres, clf_conf_thres, img_index=
                 labels = results_det["labels"]
     else:
         labels = results_det["labels"]
-    annotated_img = box_annotator.annotate(scene=input_img, detections=results_det["detections"], labels=labels)
+    annotated_img = lab_annotator.annotate(
+        scene=box_annotator.annotate(
+            scene=input_img,
+            detections=results_det["detections"],
+        ),
+        detections=results_det["detections"],
+        labels=labels,
+    )
     return annotated_img
 
 

demo/image_demo.py

@@ -26,7 +26,7 @@
 
 #%% 
 # Initializing the MegaDetectorV5 model for image detection
-detection_model = pw_detection.MegaDetectorV5(device=DEVICE, pretrained=True, verison="a")
+detection_model = pw_detection.MegaDetectorV5(device=DEVICE, pretrained=True, version="a")
 
 #%% Single image detection
 # Specifying the path to the target image TODO: Allow argparsing
@@ -45,6 +45,9 @@
 # Saving the detection results 
 pw_utils.save_detection_images(results, os.path.join(".","demo_output"), overwrite=False)
 
+# Saving the detected objects as cropped images
+pw_utils.save_crop_images(results, os.path.join(".","crop_output"), overwrite=False)
+
 #%% Batch detection
 """ Batch-detection demo """
 

demo/image_detection_demo.ipynb

@@ -64,7 +64,9 @@
    "source": [
     "# Setting the device to use for computations ('cuda' indicates GPU)\n",
     "DEVICE = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
-    "detection_model = pw_detection.MegaDetectorV5(device=DEVICE, pretrained=True, verison=\"a\")"
+    "if DEVICE == \"cuda\":\n",
+    "    torch.cuda.set_device(0)\n",
+    "detection_model = pw_detection.MegaDetectorV5(device=DEVICE, pretrained=True, version=\"a\")"
    ]
   },
   {
@@ -237,7 +239,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.8.18"
+   "version": "3.8.19"
   }
  },
  "nbformat": 4,

demo/image_separation_demo.ipynb

@@ -64,6 +64,8 @@
    "source": [
     "# Setting the device to use for computations ('cuda' indicates GPU)\n",
     "DEVICE = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
+    "if DEVICE == \"cuda\":\n",
+    "    torch.cuda.set_device(0)\n",
     "detection_model = pw_detection.MegaDetectorV5(device=DEVICE, pretrained=True)"
    ]
   },

demo/video_demo.py

@@ -41,7 +41,8 @@
 
 #%% 
 # Initializing a box annotator for visualizing detections
-box_annotator = sv.BoxAnnotator(thickness=4, text_thickness=4, text_scale=2)
+box_annotator = sv.BoundingBoxAnnotator(thickness=4)
+lab_annotator = sv.LabelAnnotator(text_color=sv.Color.BLACK, text_thickness=4, text_scale=2)
 
 def callback(frame: np.ndarray, index: int) -> np.ndarray:
     """
@@ -64,7 +65,14 @@ def callback(frame: np.ndarray, index: int) -> np.ndarray:
         results_clf = classification_model.single_image_classification(trans_clf(Image.fromarray(cropped_image)))
         labels.append("{} {:.2f}".format(results_clf["prediction"], results_clf["confidence"]))
 
-    annotated_frame = box_annotator.annotate(scene=frame, detections=results_det["detections"], labels=labels)
+    annotated_frame = lab_annotator.annotate(
+        scene=box_annotator.annotate(
+            scene=frame,
+            detections=results_det["detections"],
+        ),
+        detections=results_det["detections"],
+        labels=labels,
+    )
 
     return annotated_frame