Fix potential array overrun risk

open-mmlab · Jun 8, 2023 · 1c184d1 · 1c184d1
1 parent 6b933cc
commit 1c184d1
Showing 1 changed file with 10 additions and 6 deletions.
diff --git a/mmrotate/datasets/dota.py b/mmrotate/datasets/dota.py
@@ -232,7 +232,7 @@ def extract_xy(img_id):
             Returns:
                 Tuple of two integers, the x and y coordinates.
             """
-            pattern = re.compile(r"__(\d+)___(\d+)")
+            pattern = re.compile(r'__(\d+)___(\d+)')
             match = pattern.search(img_id)
             if match:
                 x, y = int(match.group(1)), int(match.group(2))
@@ -249,19 +249,23 @@ def extract_xy(img_id):
             for i, dets in enumerate(result):
                 bboxes, scores = dets[:, :-1], dets[:, [-1]]
                 ori_bboxes = bboxes.copy()
-                ori_bboxes[..., :2] = ori_bboxes[..., :2] + np.array([x, y], dtype=np.float32)
+                ori_bboxes[..., :2] = ori_bboxes[..., :2] + np.array(
+                    [x, y], dtype=np.float32)
                 labels = np.zeros((bboxes.shape[0], 1)) + i
-                new_result.append(np.concatenate([labels, ori_bboxes, scores], axis=1))
+                new_result.append(
+                    np.concatenate([labels, ori_bboxes, scores], axis=1))
             new_result = np.concatenate(new_result, axis=0)
             collector[oriname].append(new_result)
 
         merge_func = partial(_merge_func, CLASSES=self.CLASSES, iou_thr=0.1)
         if nproc <= 1:
             print('Executing on Single Processor')
-            merged_results = mmcv.track_iter_progress((map(merge_func, collector.items()), len(collector)))
+            merged_results = mmcv.track_iter_progress(
+                (map(merge_func, collector.items()), len(collector)))
         else:
-            print(f"Executing on {nproc} processors")
-            merged_results = mmcv.track_parallel_progress(merge_func, list(collector.items()), nproc)
+            print(f'Executing on {nproc} processors')
+            merged_results = mmcv.track_parallel_progress(
+                merge_func, list(collector.items()), nproc)
 
         # Return a zipped list of merged results
         return zip(*merged_results)