Exclude classes from inference using pretrained or custom models (#1104)

Co-authored-by: fatih akyon <[email protected]>
Co-authored-by: Copilot <[email protected]>

Author: 3 people

sahi/predict.py

@@ -57,13 +57,24 @@
 logger = logging.getLogger(__name__)
 
 
+def filter_predictions(object_prediction_list, exclude_classes_by_name, exclude_classes_by_id):
+    return [
+        obj_pred
+        for obj_pred in object_prediction_list
+        if obj_pred.category.name not in (exclude_classes_by_name or [])
+        and obj_pred.category.id not in (exclude_classes_by_id or [])
+    ]
+
+
 def get_prediction(
     image,
     detection_model,
     shift_amount: list = [0, 0],
     full_shape=None,
     postprocess: Optional[PostprocessPredictions] = None,
     verbose: int = 0,
+    exclude_classes_by_name: Optional[List[str]] = None,
+    exclude_classes_by_id: Optional[List[int]] = None,
 ) -> PredictionResult:
     """
     Function for performing prediction for given image using given detection_model.
@@ -81,7 +92,12 @@ def get_prediction(
         verbose: int
             0: no print (default)
             1: print prediction duration
-
+        exclude_classes_by_name: Optional[List[str]]
+            None: if no classes are excluded
+            List[str]: set of classes to exclude using its/their class label name/s
+        exclude_classes_by_id: Optional[List[int]]
+            None: if no classes are excluded
+            List[int]: set of classes to exclude using one or more IDs
     Returns:
         A dict with fields:
             object_prediction_list: a list of ObjectPrediction
@@ -105,6 +121,7 @@ def get_prediction(
         full_shape=full_shape,
     )
     object_prediction_list: List[ObjectPrediction] = detection_model.object_prediction_list
+    object_prediction_list = filter_predictions(object_prediction_list, exclude_classes_by_name, exclude_classes_by_id)
 
     # postprocess matching predictions
     if postprocess is not None:
@@ -142,6 +159,8 @@ def get_sliced_prediction(
     auto_slice_resolution: bool = True,
     slice_export_prefix: Optional[str] = None,
     slice_dir: Optional[str] = None,
+    exclude_classes_by_name: Optional[List[str]] = None,
+    exclude_classes_by_id: Optional[List[int]] = None,
 ) -> PredictionResult:
     """
     Function for slice image + get predicion for each slice + combine predictions in full image.
@@ -191,7 +210,12 @@ def get_sliced_prediction(
             Prefix for the exported slices. Defaults to None.
         slice_dir: str
             Directory to save the slices. Defaults to None.
-
+        exclude_classes_by_name: Optional[List[str]]
+            None: if no classes are excluded
+            List[str]: set of classes to exclude using its/their class label name/s
+        exclude_classes_by_id: Optional[List[int]]
+            None: if no classes are excluded
+            List[int]: set of classes to exclude using one or more IDs
     Returns:
         A Dict with fields:
             object_prediction_list: a list of sahi.prediction.ObjectPrediction
@@ -257,6 +281,8 @@ def get_sliced_prediction(
                 slice_image_result.original_image_height,
                 slice_image_result.original_image_width,
             ],
+            exclude_classes_by_name=exclude_classes_by_name,
+            exclude_classes_by_id=exclude_classes_by_id,
         )
         # convert sliced predictions to full predictions
         for object_prediction in prediction_result.object_prediction_list:
@@ -278,6 +304,8 @@ def get_sliced_prediction(
                 slice_image_result.original_image_width,
             ],
             postprocess=None,
+            exclude_classes_by_name=exclude_classes_by_name,
+            exclude_classes_by_id=exclude_classes_by_id,
         )
         object_prediction_list.extend(prediction_result.object_prediction_list)
 
@@ -380,6 +408,8 @@ def predict(
     verbose: int = 1,
     return_dict: bool = False,
     force_postprocess_type: bool = False,
+    exclude_classes_by_name: Optional[List[str]] = None,
+    exclude_classes_by_id: Optional[List[int]] = None,
     **kwargs,
 ):
     """
@@ -466,6 +496,12 @@ def predict(
             If True, returns a dict with 'export_dir' field.
         force_postprocess_type: bool
             If True, auto postprocess check will e disabled
+        exclude_classes_by_name: Optional[List[str]]
+            None: if no classes are excluded
+            List[str]: set of classes to exclude using its/their class label name/s
+        exclude_classes_by_id: Optional[List[int]]
+            None: if no classes are excluded
+            List[int]: set of classes to exclude using one or more IDs
     """
     # assert prediction type
     if no_standard_prediction and no_sliced_prediction:
@@ -574,6 +610,8 @@ def predict(
                 postprocess_match_threshold=postprocess_match_threshold,
                 postprocess_class_agnostic=postprocess_class_agnostic,
                 verbose=1 if verbose else 0,
+                exclude_classes_by_name=exclude_classes_by_name,
+                exclude_classes_by_id=exclude_classes_by_id,
             )
             object_prediction_list = prediction_result.object_prediction_list
             if prediction_result.durations_in_seconds:
@@ -587,6 +625,8 @@ def predict(
                 full_shape=None,
                 postprocess=None,
                 verbose=0,
+                exclude_classes_by_name=exclude_classes_by_name,
+                exclude_classes_by_id=exclude_classes_by_id,
             )
             object_prediction_list = prediction_result.object_prediction_list
 
@@ -753,6 +793,8 @@ def predict_fiftyone(
     postprocess_match_threshold: float = 0.5,
     postprocess_class_agnostic: bool = False,
     verbose: int = 1,
+    exclude_classes_by_name: Optional[List[str]] = None,
+    exclude_classes_by_id: Optional[List[int]] = None,
 ):
     """
     Performs prediction for all present images in given folder.
@@ -811,6 +853,12 @@ def predict_fiftyone(
         verbose: int
             0: no print
             1: print slice/prediction durations, number of slices, model loading/file exporting durations
+        exclude_classes_by_name: Optional[List[str]]
+            None: if no classes are excluded
+            List[str]: set of classes to exclude using its/their class label name/s
+        exclude_classes_by_id: Optional[List[int]]
+            None: if no classes are excluded
+            List[int]: set of classes to exclude using one or more IDs
     """
     check_requirements(["fiftyone"])
 
@@ -863,6 +911,8 @@ def predict_fiftyone(
                     postprocess_match_metric=postprocess_match_metric,
                     postprocess_class_agnostic=postprocess_class_agnostic,
                     verbose=verbose,
+                    exclude_classes_by_name=exclude_classes_by_name,
+                    exclude_classes_by_id=exclude_classes_by_id,
                 )
                 durations_in_seconds["slice"] += prediction_result.durations_in_seconds["slice"]
             else:
@@ -874,6 +924,8 @@ def predict_fiftyone(
                     full_shape=None,
                     postprocess=None,
                     verbose=0,
+                    exclude_classes_by_name=exclude_classes_by_name,
+                    exclude_classes_by_id=exclude_classes_by_id,
                 )
                 durations_in_seconds["prediction"] += prediction_result.durations_in_seconds["prediction"]