Mirrored_Repos/InvokeAI
1
0
Fork 0
mirror of https://github.com/invoke-ai/InvokeAI synced 2024-08-30 20:32:17 +00:00
Code Issues Packages Projects Releases Wiki Activity

(minor) Move apply_polygon_refinement condition up a layer.

Browse Source
This commit is contained in:
Ryan Dick 2024-07-31 08:50:56 -04:00
parent 33e8604b57
commit e8ecf5e155
1 changed files with 8 additions and 7 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

15
invokeai/app/invocations/grounded_sam.py
View File

@ -141,7 +141,8 @@ class GroundedSAMInvocation(BaseInvocation):
masks = sam_pipeline.segment(image=image, detection_results=detection_results)
masks = self._to_numpy_masks(masks)
masks = self._apply_polygon_refinement(masks)
if self.apply_polygon_refinement:
masks = self._apply_polygon_refinement(masks)
for detection_result, mask in zip(detection_results, masks, strict=True):
detection_result.mask = mask
@ -166,12 +167,12 @@ class GroundedSAMInvocation(BaseInvocation):
- Removes small mask pieces.
- Removes holes from the mask.
"""
if self.apply_polygon_refinement:
for idx, mask in enumerate(masks):
shape = mask.shape
polygon = mask_to_polygon(mask)
mask = polygon_to_mask(polygon, shape)
masks[idx] = mask
for idx, mask in enumerate(masks):
shape = mask.shape
assert len(shape) == 2 # Assert length to satisfy type checker.
polygon = mask_to_polygon(mask)
mask = polygon_to_mask(polygon, shape)
masks[idx] = mask
return masks