fix: 🐞 fix incorrect type annotations in postprocess module (#1327)

Author: yogendrasinghx

sahi/postprocess/combine.py

@@ -10,7 +10,7 @@
 from sahi.utils.import_utils import check_requirements
 
 
-def batched_nms(predictions: torch.tensor, match_metric: str = "IOU", match_threshold: float = 0.5):
+def batched_nms(predictions: torch.Tensor, match_metric: str = "IOU", match_threshold: float = 0.5) -> list[int]:
     """Apply non-maximum suppression to avoid detecting too many overlapping bounding boxes for a given object.
 
     Args:
@@ -40,7 +40,7 @@ def nms(
     predictions: torch.Tensor,
     match_metric: str = "IOU",
     match_threshold: float = 0.5,
-):
+) -> list[int]:
     """
     Optimized non-maximum suppression for axis-aligned bounding boxes using STRTree.
 
@@ -145,10 +145,10 @@ def nms(
 
 
 def batched_greedy_nmm(
-    object_predictions_as_tensor: torch.tensor,
+    object_predictions_as_tensor: torch.Tensor,
     match_metric: str = "IOU",
     match_threshold: float = 0.5,
-):
+) -> dict[int, list[int]]:
     """Apply greedy version of non-maximum merging per category to avoid detecting too many overlapping bounding boxes
     for a given object.
 
@@ -159,7 +159,7 @@ def batched_greedy_nmm(
         match_threshold: (float) The overlap thresh for
             match metric.
     Returns:
-        keep_to_merge_list: (Dict[int:List[int]]) mapping from prediction indices
+        keep_to_merge_list: (dict[int, list[int]]) mapping from prediction indices
         to keep to a list of prediction indices to be merged.
     """
     category_ids = object_predictions_as_tensor[:, 5].squeeze()
@@ -179,7 +179,7 @@ def greedy_nmm(
     object_predictions_as_tensor: torch.Tensor,
     match_metric: str = "IOU",
     match_threshold: float = 0.5,
-):
+) -> dict[int, list[int]]:
     """
     Optimized greedy non-maximum merging for axis-aligned bounding boxes using STRTree.
 
@@ -290,7 +290,7 @@ def batched_nmm(
     object_predictions_as_tensor: torch.Tensor,
     match_metric: str = "IOU",
     match_threshold: float = 0.5,
-):
+) -> dict[int, list[int]]:
     """Apply non-maximum merging per category to avoid detecting too many overlapping bounding boxes for a given object.
 
     Args:
@@ -300,7 +300,7 @@ def batched_nmm(
         match_threshold: (float) The overlap thresh for
             match metric.
     Returns:
-        keep_to_merge_list: (Dict[int:List[int]]) mapping from prediction indices
+        keep_to_merge_list: (dict[int, list[int]]) mapping from prediction indices
         to keep to a list of prediction indices to be merged.
     """
     category_ids = object_predictions_as_tensor[:, 5].squeeze()
@@ -320,7 +320,7 @@ def nmm(
     object_predictions_as_tensor: torch.Tensor,
     match_metric: str = "IOU",
     match_threshold: float = 0.5,
-):
+) -> dict[int, list[int]]:
     """Apply non-maximum merging to avoid detecting too many overlapping bounding boxes for a given object.
 
     Args:
@@ -329,7 +329,7 @@ def nmm(
         match_metric: (str) IOU or IOS
         match_threshold: (float) The overlap thresh for match metric.
     Returns:
-        keep_to_merge_list: (Dict[int:List[int]]) mapping from prediction indices
+        keep_to_merge_list: (dict[int, list[int]]) mapping from prediction indices
         to keep to a list of prediction indices to be merged.
     """
     # Extract coordinates and scores as tensors
@@ -460,15 +460,15 @@ def __init__(
 
         check_requirements(["torch"])
 
-    def __call__(self, predictions: list[ObjectPrediction]):
+    def __call__(self, predictions: list[ObjectPrediction]) -> list[ObjectPrediction]:
         raise NotImplementedError()
 
 
 class NMSPostprocess(PostprocessPredictions):
     def __call__(
         self,
         object_predictions: list[ObjectPrediction],
-    ):
+    ) -> list[ObjectPrediction]:
         object_prediction_list = ObjectPredictionList(object_predictions)
         object_predictions_as_torch = object_prediction_list.totensor()
         if self.class_agnostic:
@@ -491,7 +491,7 @@ class NMMPostprocess(PostprocessPredictions):
     def __call__(
         self,
         object_predictions: list[ObjectPrediction],
-    ):
+    ) -> list[ObjectPrediction]:
         object_prediction_list = ObjectPredictionList(object_predictions)
         object_predictions_as_torch = object_prediction_list.totensor()
         if self.class_agnostic:
@@ -528,7 +528,7 @@ class GreedyNMMPostprocess(PostprocessPredictions):
     def __call__(
         self,
         object_predictions: list[ObjectPrediction],
-    ):
+    ) -> list[ObjectPrediction]:
         object_prediction_list = ObjectPredictionList(object_predictions)
         object_predictions_as_torch = object_prediction_list.totensor()
         if self.class_agnostic:
@@ -566,7 +566,7 @@ class LSNMSPostprocess(PostprocessPredictions):
     def __call__(
         self,
         object_predictions: list[ObjectPrediction],
-    ):
+    ) -> list[ObjectPrediction]:
         try:
             from lsnms import nms
         except ModuleNotFoundError:

sahi/postprocess/utils.py

@@ -13,11 +13,11 @@
 
 
 class ObjectPredictionList(Sequence):
-    def __init__(self, list):
-        self.list = list
+    def __init__(self, prediction_list: list[ObjectPrediction]) -> None:
+        self.list: list[ObjectPrediction] = prediction_list
         super().__init__()
 
-    def __getitem__(self, i):
+    def __getitem__(self, i: int | tuple | list | torch.Tensor | np.ndarray) -> ObjectPredictionList:
         if torch.is_tensor(i) or isinstance(i, np.ndarray):
             i = i.tolist()
         if isinstance(i, int):
@@ -28,7 +28,11 @@ def __getitem__(self, i):
         else:
             raise NotImplementedError(f"{type(i)}")
 
-    def __setitem__(self, i, elem):
+    def __setitem__(
+        self,
+        i: int | tuple | list | torch.Tensor | np.ndarray,
+        elem: ObjectPredictionList | ObjectPrediction | list[ObjectPrediction],
+    ) -> None:
         if torch.is_tensor(i) or isinstance(i, np.ndarray):
             i = i.tolist()
         if isinstance(i, int):
@@ -45,22 +49,22 @@ def __setitem__(self, i, elem):
         else:
             raise NotImplementedError(f"{type(i)}")
 
-    def __len__(self):
+    def __len__(self) -> int:
         return len(self.list)
 
-    def __str__(self):
+    def __str__(self) -> str:
         return str(self.list)
 
-    def extend(self, object_prediction_list):
+    def extend(self, object_prediction_list: ObjectPredictionList) -> None:
         self.list.extend(object_prediction_list.list)
 
-    def totensor(self):
+    def totensor(self) -> torch.Tensor:
         return object_prediction_list_to_torch(self)
 
-    def tonumpy(self):
+    def tonumpy(self) -> np.ndarray:
         return object_prediction_list_to_numpy(self)
 
-    def tolist(self):
+    def tolist(self) -> ObjectPrediction | list[ObjectPrediction]:
         if len(self.list) == 1:
             return self.list[0]
         else:
@@ -100,7 +104,7 @@ def repair_multipolygon(shapely_multipolygon: MultiPolygon) -> MultiPolygon:
     return shapely_multipolygon
 
 
-def coco_segmentation_to_shapely(segmentation: list | list[list]):
+def coco_segmentation_to_shapely(segmentation: list | list[list]) -> MultiPolygon:
     """Fix segment data in COCO format :param segmentation: segment data in COCO format :return:"""
     if isinstance(segmentation, list) and all([not isinstance(seg, list) for seg in segmentation]):
         segmentation = [segmentation]
@@ -120,7 +124,7 @@ def coco_segmentation_to_shapely(segmentation: list | list[list]):
     return shapely_multipolygon
 
 
-def object_prediction_list_to_torch(object_prediction_list: ObjectPredictionList) -> torch.tensor:
+def object_prediction_list_to_torch(object_prediction_list: ObjectPredictionList) -> torch.Tensor:
     """
     Returns:
         torch.tensor of size N x [x1, y1, x2, y2, score, category_id]