import cv2 import numpy as np 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( mask, [contour], -1, (0 if mode == 'islands' else 255), -1 ) return np.where(mask > 0, 1, 0)