Skip to content

Commit bdae81e

Browse files
committed
(minor) Simplify GroundedSAMInvocation._filter_detections()
1 parent 67c32f3 commit bdae81e

File tree

1 file changed

+2
-8
lines changed

1 file changed

+2
-8
lines changed

invokeai/app/invocations/grounded_sam.py

Lines changed: 2 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -182,16 +182,10 @@ def _filter_detections(self, detections: list[DetectionResult]) -> list[Detectio
182182
return detections
183183
elif self.mask_filter == "largest":
184184
# Find the largest mask.
185-
mask_areas = [detection.mask.sum() for detection in detections]
186-
largest_mask_idx = mask_areas.index(max(mask_areas))
187-
return [detections[largest_mask_idx]]
185+
return [max(detections, key=lambda x: x.mask.sum())]
188186
elif self.mask_filter == "highest_box_score":
189187
# Find the detection with the highest box score.
190-
max_score_detection = detections[0]
191-
for detection in detections:
192-
if detection.score > max_score_detection.score:
193-
max_score_detection = detection
194-
return [max_score_detection]
188+
return [max(detections, key=lambda x: x.score)]
195189
else:
196190
raise ValueError(f"Invalid mask filter: {self.mask_filter}")
197191

0 commit comments

Comments
 (0)