From 1c184d1ca6698db86fa7f696554da78a447e1064 Mon Sep 17 00:00:00 2001 From: kv-chiu Date: Fri, 9 Jun 2023 02:25:09 +0800 Subject: [PATCH] Fix potential array overrun risk --- mmrotate/datasets/dota.py | 16 ++++++++++------ 1 file changed, 10 insertions(+), 6 deletions(-) diff --git a/mmrotate/datasets/dota.py b/mmrotate/datasets/dota.py index 0a8f5b08..d907e449 100644 --- 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)