Fix linting

Cubevoid · Apr 17, 2024 · b4b32e6 · b4b32e6
1 parent c058d82
commit b4b32e6
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/src/scripts/dataset_stats.py b/src/scripts/dataset_stats.py
@@ -6,14 +6,14 @@
 from src.data_collection.data_loader import DataLoader
 
 
-def calculate_iou(boxes1, boxes2):
+def calculate_iou(boxes1: npt.NDArray, boxes2: npt.NDArray) -> npt.NDArray:
     """
     Calculate IoU (Intersection over Union) between corresponding bounding boxes.
     boxes1 and boxes2 should have the format (T, O, 4), where T is the number of timesteps,
     O is the number of objects, and 4 represents (x2, y2, x1, y1).
     """
-    x1_1, y1_1, x2_1, y2_1 = np.split(boxes1, 4, axis=-1)
-    x1_2, y1_2, x2_2, y2_2 = np.split(boxes2, 4, axis=-1)
+    x1_1, y1_1, x2_1, y2_1 = np.split(boxes1, 4, axis=-1)  # pylint: disable=unbalanced-tuple-unpacking
+    x1_2, y1_2, x2_2, y2_2 = np.split(boxes2, 4, axis=-1)  # pylint: disable=unbalanced-tuple-unpacking
     x1_1, x2_1 = np.minimum(x1_1, x2_1), np.maximum(x1_1, x2_1)
     y1_1, y2_1 = np.minimum(y1_1, y2_1), np.maximum(y1_1, y2_1)
     x1_2, x2_2 = np.minimum(x1_2, x2_2), np.maximum(x1_2, x2_2)