library.mask_processing

  1from pathlib import Path
  2import cv2
  3import numpy as np
  4import matplotlib.pyplot as plt
  5
  6from scipy.spatial.distance import cdist
  7import skimage.measure
  8
  9
 10
 11
 12def load_mask(mask_path):
 13    """Load a binary mask from a PNG file.
 14    
 15    Args:
 16        mask_path: Path to the mask file (str or Path object).
 17        
 18    Returns:
 19        Binary mask as numpy array with values 0 (background) and 255 (foreground).
 20        
 21    Raises:
 22        FileNotFoundError: If the mask file cannot be loaded.
 23    """
 24    mask = cv2.imread(str(mask_path), cv2.IMREAD_GRAYSCALE)
 25    if mask is None:
 26        raise FileNotFoundError(f"Could not load mask from {mask_path}")
 27    return mask
 28
 29def repair_root_mask_edges(mask_path, 
 30                          min_component_size=20,
 31                          max_edge_distance=50,
 32                          max_iterations=10,
 33                          closing_kernel_height=5,
 34                          closing_kernel_width=3,
 35                          closing_iterations=2,
 36                          save_path=None,
 37                          visualize=False):
 38    """Repair root mask gaps using edge-to-edge distance measurements.
 39    
 40    Uses iterative edge-based gap filling to connect fragmented root components.
 41    Measures minimum distance between component edges rather than centroids,
 42    which is more accurate for irregular shapes like roots.
 43    
 44    Args:
 45        mask_path: Path to root mask file (str or Path).
 46        min_component_size: Minimum component area in pixels to consider for connection.
 47        max_edge_distance: Maximum pixel distance between edges to connect components.
 48        max_iterations: Maximum number of gap-filling iterations.
 49        closing_kernel_height: Height of initial morphological closing kernel.
 50        closing_kernel_width: Width of initial morphological closing kernel.
 51        closing_iterations: Number of initial closing iterations.
 52        save_path: Optional path to save repaired mask. If None, does not save.
 53        visualize: If True, displays before/after visualization with original image.
 54        
 55    Returns:
 56        Binary mask array (0 and 255) with repaired gaps.
 57        
 58    Example:
 59        repaired_mask = repair_root_mask_edges(
 60            'root_mask.png',
 61            max_edge_distance=50,
 62            visualize=True
 63        )
 64    """
 65    
 66    # Load mask
 67    mask = load_mask(mask_path)
 68    
 69    # Initial morphological closing
 70    kernel = np.ones((closing_kernel_height, closing_kernel_width), np.uint8)
 71    closed_mask = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, kernel, 
 72                                   iterations=closing_iterations)
 73    
 74    repaired_mask = closed_mask.copy()
 75    
 76    # Iterative edge-based gap filling
 77    for iteration in range(max_iterations):
 78        retval, labels = cv2.connectedComponents(repaired_mask)
 79        regions = skimage.measure.regionprops(labels)
 80        
 81        # Get substantial components with their edge pixels
 82        components = []
 83        for region in regions:
 84            if region.area < min_component_size:
 85                continue
 86            
 87            coords = region.coords
 88            min_y = coords[:, 0].min()
 89            max_y = coords[:, 0].max()
 90            
 91            # Get bottom and top edge pixels
 92            bottom_edge = coords[coords[:, 0] == max_y]
 93            top_edge = coords[coords[:, 0] == min_y]
 94            
 95            components.append({
 96                'label': region.label,
 97                'area': region.area,
 98                'min_y': min_y,
 99                'max_y': max_y,
100                'bottom_edge': bottom_edge,
101                'top_edge': top_edge
102            })
103        
104        connections_made = 0
105        
106        # Find and connect close component edges
107        for i, comp1 in enumerate(components):
108            for j, comp2 in enumerate(components):
109                if i >= j:
110                    continue
111                
112                # Check if comp2 is below comp1
113                if comp1['max_y'] >= comp2['min_y']:
114                    continue
115                
116                # Calculate minimum distance between edges
117                distances = cdist(comp1['bottom_edge'], comp2['top_edge'])
118                min_distance = distances.min()
119                min_idx = np.unravel_index(distances.argmin(), distances.shape)
120                
121                # Get the closest points
122                closest_bottom = comp1['bottom_edge'][min_idx[0]]
123                closest_top = comp2['top_edge'][min_idx[1]]
124                
125                if min_distance <= max_edge_distance:
126                    # Draw line between closest points
127                    cv2.line(repaired_mask, 
128                            (closest_bottom[1], closest_bottom[0]),
129                            (closest_top[1], closest_top[0]), 
130                            255, 3)
131                    connections_made += 1
132        
133        # Apply closing to smooth connections
134        kernel_smooth = np.ones((5, 3), np.uint8)
135        repaired_mask = cv2.morphologyEx(repaired_mask, cv2.MORPH_CLOSE, 
136                                        kernel_smooth, iterations=1)
137        
138        # Stop if no connections were made
139        if connections_made == 0:
140            break
141    
142    # Save if path provided
143    if save_path is not None:
144        cv2.imwrite(str(save_path), repaired_mask)
145    
146    # Visualize if requested
147    if visualize:
148        mask_path_obj = Path(mask_path)
149        mask_name = mask_path_obj.stem
150        image_num = mask_name.replace('test_image_', '').replace('_root', '')
151        
152        # Try to find original image
153        images_path = mask_path_obj.parent.parent.parent / 'Kaggle'
154        original_path = images_path / f'test_image_{image_num}.png'
155        
156        retval_before = cv2.connectedComponents(closed_mask)[0]
157        retval_after, labels_after = cv2.connectedComponents(repaired_mask)
158        
159        fig, axes = plt.subplots(2, 1, figsize=(15, 30))
160        
161        axes[0].imshow(labels_after, cmap='gist_ncar')
162        axes[0].set_title(f'After Edge-based Repair: {retval_after - 1} components\n'
163                         f'(started with {retval_before - 1})', fontsize=16)
164        axes[0].axis('off')
165        
166        if original_path.exists():
167            original = cv2.imread(str(original_path))
168            original_rgb = cv2.cvtColor(original, cv2.COLOR_BGR2RGB)
169            axes[1].imshow(original_rgb)
170            axes[1].imshow(labels_after, cmap='gist_ncar', alpha=0.4)
171            axes[1].set_title('Overlay with Original Image', fontsize=16)
172        else:
173            axes[1].imshow(labels_after, cmap='gist_ncar')
174            axes[1].set_title('Repaired Mask (no original image found)', fontsize=16)
175        
176        axes[1].axis('off')
177        
178        plt.tight_layout()
179        plt.show()
180    
181    return repaired_mask