zjt/SAMapi
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Move batched NMS indices to correct device (closes #17)

Browse Source
This commit is contained in:
Louis Maddox
2023-04-10 11:58:44 -07:00
committed by ericmintun
parent 2780a301de
commit 6325eb8512
1 changed files with 2 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files
segment_anything/automatic_mask_generator.py
+2 -2
View File
@@ -214,7 +214,7 @@ class SamAutomaticMaskGenerator:
keep_by_nms = batched_nms( keep_by_nms = batched_nms(
data["boxes"].float(), data["boxes"].float(),
scores, scores,
torch.zeros(len(data["boxes"])), # categories torch.zeros_like(data["boxes"][:,0]), # categories
iou_threshold=self.crop_nms_thresh, iou_threshold=self.crop_nms_thresh,
) )
data.filter(keep_by_nms) data.filter(keep_by_nms)
@@ -251,7 +251,7 @@ class SamAutomaticMaskGenerator:
keep_by_nms = batched_nms( keep_by_nms = batched_nms(
data["boxes"].float(), data["boxes"].float(),
data["iou_preds"], data["iou_preds"],
torch.zeros(len(data["boxes"])), # categories torch.zeros_like(data["boxes"][:,0]), # categories
iou_threshold=self.box_nms_thresh, iou_threshold=self.box_nms_thresh,
) )
data.filter(keep_by_nms) data.filter(keep_by_nms)