Perform PostProcessing on output masks.
(self, masks: torch.Tensor, orig_hw)
| 72 | return boxes |
| 73 | |
| 74 | def postprocess_masks(self, masks: torch.Tensor, orig_hw) -> torch.Tensor: |
| 75 | """ |
| 76 | Perform PostProcessing on output masks. |
| 77 | """ |
| 78 | from sam2_train.utils.misc import get_connected_components |
| 79 | |
| 80 | masks = masks.float() |
| 81 | if self.max_hole_area > 0: |
| 82 | # Holes are those connected components in background with area <= self.fill_hole_area |
| 83 | # (background regions are those with mask scores <= self.mask_threshold) |
| 84 | mask_flat = masks.flatten(0, 1).unsqueeze(1) # flatten as 1-channel image |
| 85 | labels, areas = get_connected_components(mask_flat <= self.mask_threshold) |
| 86 | is_hole = (labels > 0) & (areas <= self.max_hole_area) |
| 87 | is_hole = is_hole.reshape_as(masks) |
| 88 | # We fill holes with a small positive mask score (10.0) to change them to foreground. |
| 89 | masks = torch.where(is_hole, self.mask_threshold + 10.0, masks) |
| 90 | |
| 91 | if self.max_sprinkle_area > 0: |
| 92 | labels, areas = get_connected_components(mask_flat > self.mask_threshold) |
| 93 | is_hole = (labels > 0) & (areas <= self.max_sprinkle_area) |
| 94 | is_hole = is_hole.reshape_as(masks) |
| 95 | # We fill holes with negative mask score (-10.0) to change them to background. |
| 96 | masks = torch.where(is_hole, self.mask_threshold - 10.0, masks) |
| 97 | |
| 98 | masks = F.interpolate(masks, orig_hw, mode="bilinear", align_corners=False) |
| 99 | return masks |
no test coverage detected