Spaces:
Runtime error
Runtime error
import cv2 | |
import som | |
import numpy as np | |
import supervision as sv | |
class Visualizer: | |
def __init__( | |
self, | |
line_thickness: int = 2, | |
mask_opacity: float = 0.1, | |
text_scale: float = 0.6 | |
) -> None: | |
self.box_annotator = sv.BoundingBoxAnnotator( | |
color_lookup=sv.ColorLookup.INDEX, | |
thickness=line_thickness) | |
self.mask_annotator = sv.MaskAnnotator( | |
color_lookup=sv.ColorLookup.INDEX, | |
opacity=mask_opacity) | |
self.polygon_annotator = sv.PolygonAnnotator( | |
color_lookup=sv.ColorLookup.INDEX, | |
thickness=line_thickness) | |
self.label_annotator = sv.LabelAnnotator( | |
color=sv.Color.black(), | |
text_color=sv.Color.white(), | |
color_lookup=sv.ColorLookup.INDEX, | |
text_position=sv.Position.CENTER_OF_MASS, | |
text_scale=text_scale) | |
def visualize( | |
self, | |
image: np.ndarray, | |
detections: sv.Detections, | |
with_box: bool, | |
with_mask: bool, | |
with_polygon: bool, | |
with_label: bool | |
) -> np.ndarray: | |
annotated_image = image.copy() | |
if with_box: | |
annotated_image = self.box_annotator.annotate( | |
scene=annotated_image, detections=detections) | |
if with_mask: | |
annotated_image = self.mask_annotator.annotate( | |
scene=annotated_image, detections=detections) | |
if with_polygon: | |
annotated_image = self.polygon_annotator.annotate( | |
scene=annotated_image, detections=detections) | |
if with_label: | |
labels = list(map(str, range(len(detections)))) | |
annotated_image = self.label_annotator.annotate( | |
scene=annotated_image, detections=detections, labels=labels) | |
return annotated_image | |
def refine_mask( | |
mask: np.ndarray, | |
area_threshold: float, | |
mode: str = 'islands' | |
) -> np.ndarray: | |
""" | |
Refines a mask by removing small islands or filling small holes based on area | |
threshold. | |
Parameters: | |
mask (np.ndarray): Input binary mask. | |
area_threshold (float): Threshold for relative area to remove or fill features. | |
mode (str): Operation mode ('islands' for removing islands, 'holes' for filling | |
holes). | |
Returns: | |
np.ndarray: Refined binary mask. | |
""" | |
mask = np.uint8(mask * 255) | |
operation = cv2.RETR_EXTERNAL if mode == 'islands' else cv2.RETR_CCOMP | |
contours, _ = cv2.findContours( | |
mask, operation, cv2.CHAIN_APPROX_SIMPLE | |
) | |
total_area = cv2.countNonZero(mask) if mode == 'islands' else mask.size | |
for contour in contours: | |
area = cv2.contourArea(contour) | |
relative_area = area / total_area | |
if relative_area < area_threshold: | |
cv2.drawContours( | |
image=mask, | |
contours=[contour], | |
contourIdx=-1, | |
color=(0 if mode == 'islands' else 255), | |
thickness=-1 | |
) | |
return np.where(mask > 0, 1, 0).astype(bool) | |
def filter_masks_by_relative_area( | |
masks: np.ndarray, | |
min_relative_area: float = 0.02, | |
max_relative_area: float = 1.0 | |
) -> np.ndarray: | |
""" | |
Filters out masks based on their relative area. | |
Parameters: | |
masks (np.ndarray): A 3D numpy array where each slice along the third dimension | |
represents a mask. | |
min_relative_area (float): Minimum relative area threshold for keeping a mask. | |
max_relative_area (float): Maximum relative area threshold for keeping a mask. | |
Returns: | |
np.ndarray: A 3D numpy array of filtered masks. | |
""" | |
mask_areas = masks.sum(axis=(1, 2)) | |
total_area = masks.shape[1] * masks.shape[2] | |
relative_areas = mask_areas / total_area | |
min_area_filter = relative_areas >= min_relative_area | |
max_area_filter = relative_areas <= max_relative_area | |
return masks[min_area_filter & max_area_filter] | |
def postprocess_masks( | |
detections: sv.Detections, | |
area_threshold: float = 0.01, | |
min_relative_area: float = 0.01, | |
max_relative_area: float = 1.0, | |
iou_threshold: float = 0.9 | |
) -> sv.Detections: | |
""" | |
Post-processes the masks of detection objects by removing small islands and filling | |
small holes. | |
Parameters: | |
detections (sv.Detections): Detection objects to be filtered. | |
area_threshold (float): Threshold for relative area to remove or fill features. | |
min_relative_area (float): Minimum relative area threshold for detections. | |
max_relative_area (float): Maximum relative area threshold for detections. | |
iou_threshold (float): The IoU threshold above which masks will be considered as | |
overlapping. | |
Returns: | |
np.ndarray: Post-processed masks. | |
""" | |
masks = detections.mask.copy() | |
for i in range(len(masks)): | |
masks[i] = refine_mask( | |
mask=masks[i], | |
area_threshold=area_threshold, | |
mode='islands' | |
) | |
masks[i] = refine_mask( | |
mask=masks[i], | |
area_threshold=area_threshold, | |
mode='holes' | |
) | |
masks = filter_masks_by_relative_area( | |
masks=masks, | |
min_relative_area=min_relative_area, | |
max_relative_area=max_relative_area) | |
masks = som.mask_non_max_suppression( | |
masks=masks, | |
iou_threshold=iou_threshold) | |
return sv.Detections( | |
xyxy=sv.mask_to_xyxy(masks), | |
mask=masks | |
) | |