Perform PostProcessing on output masks.
(self, masks: torch.Tensor, orig_hw)
| 68 | return boxes |
| 69 | |
| 70 | def postprocess_masks(self, masks: torch.Tensor, orig_hw) -> torch.Tensor: |
| 71 | """ |
| 72 | Perform PostProcessing on output masks. |
| 73 | """ |
| 74 | from sam2.utils.misc import get_connected_components |
| 75 | |
| 76 | masks = masks.float() |
| 77 | input_masks = masks |
| 78 | mask_flat = masks.flatten(0, 1).unsqueeze(1) # flatten as 1-channel image |
| 79 | try: |
| 80 | if self.max_hole_area > 0: |
| 81 | # Holes are those connected components in background with area <= self.fill_hole_area |
| 82 | # (background regions are those with mask scores <= self.mask_threshold) |
| 83 | labels, areas = get_connected_components(mask_flat <= self.mask_threshold) |
| 84 | is_hole = (labels > 0) & (areas <= self.max_hole_area) |
| 85 | is_hole = is_hole.reshape_as(masks) |
| 86 | # We fill holes with a small positive mask score (10.0) to change them to foreground. |
| 87 | masks = torch.where(is_hole, self.mask_threshold + 10.0, masks) |
| 88 | |
| 89 | if self.max_sprinkle_area > 0: |
| 90 | labels, areas = get_connected_components(mask_flat > self.mask_threshold) |
| 91 | is_hole = (labels > 0) & (areas <= self.max_sprinkle_area) |
| 92 | is_hole = is_hole.reshape_as(masks) |
| 93 | # We fill holes with negative mask score (-10.0) to change them to background. |
| 94 | masks = torch.where(is_hole, self.mask_threshold - 10.0, masks) |
| 95 | except Exception as e: |
| 96 | # Skip the post-processing step if the CUDA kernel fails |
| 97 | warnings.warn( |
| 98 | f"{e}\n\nSkipping the post-processing step due to the error above. You can " |
| 99 | "still use SAM 2 and it's OK to ignore the error above, although some post-processing " |
| 100 | "functionality may be limited (which doesn't affect the results in most cases; see " |
| 101 | "https://github.com/facebookresearch/sam2/blob/main/INSTALL.md).", |
| 102 | category=UserWarning, |
| 103 | stacklevel=2, |
| 104 | ) |
| 105 | masks = input_masks |
| 106 | |
| 107 | masks = F.interpolate(masks.float(), orig_hw, mode="bilinear", align_corners=False).to(masks.dtype) |
| 108 | return masks |
no test coverage detected