Skip to content

GitLab

Commit 6325eb85 authored by Louis Maddox's avatar Louis Maddox Committed by ericmintun
Browse files

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

parent 2780a301
Show whitespace changes
Inline Side-by-side
Showing with 2 additions and 2 deletions
segment_anything/automatic_mask_generator.py
View file @ 6325eb85
...@@ -214,7 +214,7 @@ class SamAutomaticMaskGenerator: ...@@ -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: ...@@ -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)
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment