Create clean_refine.py

clean_refine.py

@@ -0,0 +1,71 @@
+import cv2
+import numpy as np
+
+def apply_gaussian_blur(mask, kernel_size=5):
+    """Apply Gaussian blur to smooth the mask edges."""
+    return cv2.GaussianBlur(mask, (kernel_size, kernel_size), 0)
+
+def apply_threshold(mask, threshold=127):
+    """Apply binary threshold to sharpen the mask."""
+    _, binary_mask = cv2.threshold(mask, threshold, 255, cv2.THRESH_BINARY)
+    return binary_mask
+
+def refine_edges(mask, kernel_size=3):
+    """Refine edges using morphological operations."""
+    kernel = np.ones((kernel_size, kernel_size), np.uint8)
+    eroded = cv2.erode(mask, kernel, iterations=1)
+    dilated = cv2.dilate(mask, kernel, iterations=1)
+    refined = dilated - eroded
+    return cv2.bitwise_or(eroded, refined)
+
+def apply_contour_smoothing(mask):
+    """Smooth contours of the mask."""
+    contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
+    smooth_mask = np.zeros_like(mask)
+    for contour in contours:
+        epsilon = 0.02 * cv2.arcLength(contour, True)
+        approx = cv2.approxPolyDP(contour, epsilon, True)
+        cv2.drawContours(smooth_mask, [approx], 0, 255, -1)
+    return smooth_mask
+
+def refine_mask(mask, blur_kernel=5, edge_kernel=3, threshold_value=127):
+    """Apply a series of refinement operations to the mask."""
+    mask = apply_gaussian_blur(mask, blur_kernel)
+    mask = apply_threshold(mask, threshold_value)
+    mask = refine_edges(mask, edge_kernel)
+    mask = apply_contour_smoothing(mask)
+    return mask
+    
+
+def apply_morphology(mask, kernel_size=3):
+    """Apply morphological operations to clean up the mask."""
+    kernel = np.ones((kernel_size, kernel_size), np.uint8)
+    
+    # Opening operation to remove small noise
+    mask = cv2.morphologyEx(mask, cv2.MORPH_OPEN, kernel)
+    
+    # Closing operation to fill small holes
+    mask = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, kernel)
+    
+    return mask
+
+def remove_small_objects(mask, min_size=100):
+    """Remove small objects from the mask based on area."""
+    # Ensure mask is binary and single channel
+    if len(mask.shape) > 2:
+        mask = cv2.cvtColor(mask, cv2.COLOR_BGR2GRAY)
+    _, binary_mask = cv2.threshold(mask, 127, 255, cv2.THRESH_BINARY)
+    
+    num_labels, labels, stats, _ = cv2.connectedComponentsWithStats(binary_mask, connectivity=8)
+    
+    for i in range(1, num_labels):  # Start from 1 to skip the background
+        if stats[i, cv2.CC_STAT_AREA] < min_size:
+            binary_mask[labels == i] = 0
+    
+    return binary_mask
+
+def clean_mask(mask, morph_kernel_size=3, min_object_size=100):
+    """Apply both morphological operations and small object removal."""
+    mask = apply_morphology(mask, kernel_size=morph_kernel_size)
+    mask = remove_small_objects(mask, min_size=min_object_size)
+    return mask