Skip to content

GitLab

Commit 7fa17d78 authored by Eric Mintun's avatar Eric Mintun
Browse files

Lint.

parent b028d543
Hide whitespace changes
Inline Side-by-side
Showing with 3 additions and 3 deletions
segment_anything/automatic_mask_generator.py
View file @ 7fa17d78
...@@ -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_like(data["boxes"][:,0]), # 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_like(data["boxes"][:,0]), # 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)
...@@ -357,7 +357,7 @@ class SamAutomaticMaskGenerator: ...@@ -357,7 +357,7 @@ class SamAutomaticMaskGenerator:
keep_by_nms = batched_nms( keep_by_nms = batched_nms(
boxes.float(), boxes.float(),
torch.as_tensor(scores), torch.as_tensor(scores),
torch.zeros_like(boxes[:,0]), # categories torch.zeros_like(boxes[:, 0]), # categories
iou_threshold=nms_thresh, iou_threshold=nms_thresh,
) )
......
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