From 3cb96eaf1caaf3589556b5703a11f040fb832c40 Mon Sep 17 00:00:00 2001 From: javierganan99 Date: Mon, 16 Sep 2024 10:45:47 +0200 Subject: [PATCH 1/2] Returning the indices sorted by IoU in descending order --- fastsam/utils.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/fastsam/utils.py b/fastsam/utils.py index 33d37cd..7b550e5 100644 --- a/fastsam/utils.py +++ b/fastsam/utils.py @@ -44,7 +44,7 @@ def bbox_iou(box1, boxes, iou_thres=0.9, image_shape=(640, 640), raw_output=Fals box1: (4, ) boxes: (n, 4) Returns: - high_iou_indices: Indices of boxes with IoU > thres + high_iou_indices: Indices of boxes with IoU > thres sorted in descending order ''' boxes = adjust_bboxes_to_image_border(boxes, image_shape) # obtain coordinates for intersections @@ -72,8 +72,9 @@ def bbox_iou(box1, boxes, iou_thres=0.9, image_shape=(640, 640), raw_output=Fals # get indices of boxes with IoU > thres high_iou_indices = torch.nonzero(iou > iou_thres).flatten() + sorted_high_iou_indices = high_iou_indices[torch.argsort(iou[high_iou_indices], descending=True)] - return high_iou_indices + return sorted_high_iou_indices def image_to_np_ndarray(image): From 9e9899d9cb31707b1a3a46d2610d198dfb8a646d Mon Sep 17 00:00:00 2001 From: javierganan99 Date: Mon, 16 Sep 2024 10:46:12 +0200 Subject: [PATCH 2/2] Taking the first IoU index to avoid Exception --- fastsam/predict.py | 1 + 1 file changed, 1 insertion(+) diff --git a/fastsam/predict.py b/fastsam/predict.py index cc15128..727a57d 100644 --- a/fastsam/predict.py +++ b/fastsam/predict.py @@ -31,6 +31,7 @@ def postprocess(self, preds, img, orig_imgs): full_box = full_box.view(1, -1) critical_iou_index = bbox_iou(full_box[0][:4], p[0][:, :4], iou_thres=0.9, image_shape=img.shape[2:]) if critical_iou_index.numel() != 0: + critical_iou_index = critical_iou_index[:1] full_box[0][4] = p[0][critical_iou_index][:,4] full_box[0][6:] = p[0][critical_iou_index][:,6:] p[0][critical_iou_index] = full_box