Roboflow/SoM

1import cv22import som3 4import numpy as np5import supervision as sv6 7 8class Visualizer:9 10    def __init__(11        self,12        line_thickness: int = 2,13        mask_opacity: float = 0.1,14        text_scale: float = 0.615    ) -> None:16        self.box_annotator = sv.BoundingBoxAnnotator(17            color_lookup=sv.ColorLookup.INDEX,18            thickness=line_thickness)19        self.mask_annotator = sv.MaskAnnotator(20            color_lookup=sv.ColorLookup.INDEX,21            opacity=mask_opacity)22        self.polygon_annotator = sv.PolygonAnnotator(23            color_lookup=sv.ColorLookup.INDEX,24            thickness=line_thickness)25        self.label_annotator = sv.LabelAnnotator(26            color=sv.Color.black(),27            text_color=sv.Color.white(),28            color_lookup=sv.ColorLookup.INDEX,29            text_position=sv.Position.CENTER_OF_MASS,30            text_scale=text_scale)31 32    def visualize(33        self,34        image: np.ndarray,35        detections: sv.Detections,36        with_box: bool,37        with_mask: bool,38        with_polygon: bool,39        with_label: bool40    ) -> np.ndarray:41        annotated_image = image.copy()42        if with_box:43            annotated_image = self.box_annotator.annotate(44                scene=annotated_image, detections=detections)45        if with_mask:46            annotated_image = self.mask_annotator.annotate(47                scene=annotated_image, detections=detections)48        if with_polygon:49            annotated_image = self.polygon_annotator.annotate(50                scene=annotated_image, detections=detections)51        if with_label:52            labels = list(map(str, range(len(detections))))53            annotated_image = self.label_annotator.annotate(54                scene=annotated_image, detections=detections, labels=labels)55        return annotated_image56 57 58def refine_mask(59    mask: np.ndarray,60    area_threshold: float,61    mode: str = 'islands'62) -> np.ndarray:63    """64    Refines a mask by removing small islands or filling small holes based on area65    threshold.66 67    Parameters:68        mask (np.ndarray): Input binary mask.69        area_threshold (float): Threshold for relative area to remove or fill features.70        mode (str): Operation mode ('islands' for removing islands, 'holes' for filling71                    holes).72 73    Returns:74        np.ndarray: Refined binary mask.75    """76    mask = np.uint8(mask * 255)77    operation = cv2.RETR_EXTERNAL if mode == 'islands' else cv2.RETR_CCOMP78    contours, _ = cv2.findContours(79        mask, operation, cv2.CHAIN_APPROX_SIMPLE80    )81    total_area = cv2.countNonZero(mask) if mode == 'islands' else mask.size82 83    for contour in contours:84        area = cv2.contourArea(contour)85        relative_area = area / total_area86        if relative_area < area_threshold:87            cv2.drawContours(88                image=mask,89                contours=[contour],90                contourIdx=-1,91                color=(0 if mode == 'islands' else 255),92                thickness=-193            )94 95    return np.where(mask > 0, 1, 0).astype(bool)96 97 98def filter_masks_by_relative_area(99    masks: np.ndarray,100    min_relative_area: float = 0.02,101    max_relative_area: float = 1.0102) -> np.ndarray:103    """104    Filters out masks based on their relative area.105 106    Parameters:107        masks (np.ndarray): A 3D numpy array where each slice along the third dimension108            represents a mask.109        min_relative_area (float): Minimum relative area threshold for keeping a mask.110        max_relative_area (float): Maximum relative area threshold for keeping a mask.111 112    Returns:113        np.ndarray: A 3D numpy array of filtered masks.114    """115    mask_areas = masks.sum(axis=(1, 2))116    total_area = masks.shape[1] * masks.shape[2]117    relative_areas = mask_areas / total_area118    min_area_filter = relative_areas >= min_relative_area119    max_area_filter = relative_areas <= max_relative_area120    return masks[min_area_filter & max_area_filter]121 122 123def postprocess_masks(124    detections: sv.Detections,125    area_threshold: float = 0.01,126    min_relative_area: float = 0.01,127    max_relative_area: float = 1.0,128    iou_threshold: float = 0.9129) -> sv.Detections:130    """131    Post-processes the masks of detection objects by removing small islands and filling132    small holes.133 134    Parameters:135        detections (sv.Detections): Detection objects to be filtered.136        area_threshold (float): Threshold for relative area to remove or fill features.137        min_relative_area (float): Minimum relative area threshold for detections.138        max_relative_area (float): Maximum relative area threshold for detections.139        iou_threshold (float): The IoU threshold above which masks will be considered as140            overlapping.141 142    Returns:143        np.ndarray: Post-processed masks.144    """145    masks = detections.mask.copy()146    for i in range(len(masks)):147        masks[i] = refine_mask(148            mask=masks[i],149            area_threshold=area_threshold,150            mode='islands'151        )152        masks[i] = refine_mask(153            mask=masks[i],154            area_threshold=area_threshold,155            mode='holes'156        )157    masks = filter_masks_by_relative_area(158        masks=masks,159        min_relative_area=min_relative_area,160        max_relative_area=max_relative_area)161    masks = som.mask_non_max_suppression(162        masks=masks,163        iou_threshold=iou_threshold)164 165    return sv.Detections(166        xyxy=sv.mask_to_xyxy(masks),167        mask=masks168    )169