Skip to content

Commit

Permalink
fix IDEA-Research#17 from SAM
Browse files Browse the repository at this point in the history
  • Loading branch information
hkchengrex committed Aug 28, 2023
1 parent f34e1a3 commit b4459c4
Showing 1 changed file with 3 additions and 3 deletions.
6 changes: 3 additions & 3 deletions segment_anything/segment_anything/automatic_mask_generator.py
Original file line number Diff line number Diff line change
@@ -214,7 +214,7 @@ def _generate_masks(self, image: np.ndarray) -> MaskData:
keep_by_nms = batched_nms(
data["boxes"].float(),
scores,
torch.zeros(len(data["boxes"])), # categories
torch.zeros_like(len(data["boxes"][:,0])), # categories
iou_threshold=self.crop_nms_thresh,
)
data.filter(keep_by_nms)
@@ -251,7 +251,7 @@ def _process_crop(
keep_by_nms = batched_nms(
data["boxes"].float(),
data["iou_preds"],
torch.zeros(len(data["boxes"])), # categories
torch.zeros_like(len(data["boxes"][:,0])), # categories
iou_threshold=self.box_nms_thresh,
)
data.filter(keep_by_nms)
@@ -357,7 +357,7 @@ def postprocess_small_regions(
keep_by_nms = batched_nms(
boxes.float(),
torch.as_tensor(scores),
torch.zeros(len(boxes)), # categories
torch.zeros_like(len(boxes[:,0])), # categories
iou_threshold=nms_thresh,
)

0 comments on commit b4459c4

Please sign in to comment.