geoai-py 0.17.0__py2.py3-none-any.whl → 0.18.1__py2.py3-none-any.whl

geoai/tools/multiclean.py

@@ -0,0 +1,357 @@
+"""
+MultiClean integration utilities for cleaning segmentation results.
+
+This module provides functions to use MultiClean (https://github.com/DPIRD-DMA/MultiClean)
+for post-processing segmentation masks and classification rasters. MultiClean performs
+morphological operations to smooth edges, remove noise islands, and fill gaps.
+"""
+
+import os
+from typing import Optional, List, Union, Tuple
+import numpy as np
+
+try:
+    from multiclean import clean_array
+
+    MULTICLEAN_AVAILABLE = True
+except ImportError:
+    MULTICLEAN_AVAILABLE = False
+
+try:
+    import rasterio
+
+    RASTERIO_AVAILABLE = True
+except ImportError:
+    RASTERIO_AVAILABLE = False
+
+
+def check_multiclean_available():
+    """
+    Check if multiclean is installed.
+
+    Raises:
+        ImportError: If multiclean is not installed.
+    """
+    if not MULTICLEAN_AVAILABLE:
+        raise ImportError(
+            "multiclean is not installed. "
+            "Please install it with: pip install multiclean "
+            "or: pip install geoai-py[extra]"
+        )
+
+
+def clean_segmentation_mask(
+    mask: np.ndarray,
+    class_values: Optional[Union[int, List[int]]] = None,
+    smooth_edge_size: int = 2,
+    min_island_size: int = 100,
+    connectivity: int = 8,
+    max_workers: Optional[int] = None,
+    fill_nan: bool = False,
+) -> np.ndarray:
+    """
+    Clean a segmentation mask using MultiClean morphological operations.
+
+    This function applies three cleaning operations:
+    1. Edge smoothing - Uses morphological opening to reduce jagged boundaries
+    2. Island removal - Eliminates small connected components (noise)
+    3. Gap filling - Replaces invalid pixels with nearest valid class
+
+    Args:
+        mask (np.ndarray): 2D numpy array containing segmentation classes.
+            Can be int or float. NaN values are treated as nodata.
+        class_values (int, list of int, or None): Target class values to process.
+            If None, auto-detects unique values from the mask. Defaults to None.
+        smooth_edge_size (int): Kernel width in pixels for edge smoothing.
+            Set to 0 to disable smoothing. Defaults to 2.
+        min_island_size (int): Minimum area (in pixels) for connected components.
+            Components with area strictly less than this are removed. Defaults to 100.
+        connectivity (int): Connectivity for component detection. Use 4 or 8.
+            8-connectivity considers diagonal neighbors. Defaults to 8.
+        max_workers (int, optional): Thread pool size for parallel processing.
+            If None, uses default threading. Defaults to None.
+        fill_nan (bool): Whether to fill NaN pixels with nearest valid class.
+            Defaults to False.
+
+    Returns:
+        np.ndarray: Cleaned 2D segmentation mask with same shape as input.
+
+    Raises:
+        ImportError: If multiclean is not installed.
+        ValueError: If mask is not 2D or if connectivity is not 4 or 8.
+
+    Example:
+        >>> import numpy as np
+        >>> from geoai.tools.multiclean import clean_segmentation_mask
+        >>> mask = np.random.randint(0, 3, (512, 512))
+        >>> cleaned = clean_segmentation_mask(
+        ...     mask,
+        ...     class_values=[0, 1, 2],
+        ...     smooth_edge_size=2,
+        ...     min_island_size=50
+        ... )
+    """
+    check_multiclean_available()
+
+    if mask.ndim != 2:
+        raise ValueError(f"Mask must be 2D, got shape {mask.shape}")
+
+    if connectivity not in [4, 8]:
+        raise ValueError(f"Connectivity must be 4 or 8, got {connectivity}")
+
+    # Apply MultiClean
+    cleaned = clean_array(
+        mask,
+        class_values=class_values,
+        smooth_edge_size=smooth_edge_size,
+        min_island_size=min_island_size,
+        connectivity=connectivity,
+        max_workers=max_workers,
+        fill_nan=fill_nan,
+    )
+
+    return cleaned
+
+
+def clean_raster(
+    input_path: str,
+    output_path: str,
+    class_values: Optional[Union[int, List[int]]] = None,
+    smooth_edge_size: int = 2,
+    min_island_size: int = 100,
+    connectivity: int = 8,
+    max_workers: Optional[int] = None,
+    fill_nan: bool = False,
+    band: int = 1,
+    nodata: Optional[float] = None,
+) -> None:
+    """
+    Clean a classification raster (GeoTIFF) and save the result.
+
+    Reads a GeoTIFF file, applies MultiClean morphological operations,
+    and saves the cleaned result while preserving geospatial metadata
+    (CRS, transform, nodata value).
+
+    Args:
+        input_path (str): Path to input GeoTIFF file.
+        output_path (str): Path to save cleaned GeoTIFF file.
+        class_values (int, list of int, or None): Target class values to process.
+            If None, auto-detects unique values. Defaults to None.
+        smooth_edge_size (int): Kernel width in pixels for edge smoothing.
+            Defaults to 2.
+        min_island_size (int): Minimum area (in pixels) for components.
+            Defaults to 100.
+        connectivity (int): Connectivity for component detection (4 or 8).
+            Defaults to 8.
+        max_workers (int, optional): Thread pool size. Defaults to None.
+        fill_nan (bool): Whether to fill NaN/nodata pixels. Defaults to False.
+        band (int): Band index to read (1-indexed). Defaults to 1.
+        nodata (float, optional): Nodata value to use. If None, uses value
+            from input file. Defaults to None.
+
+    Returns:
+        None: Writes cleaned raster to output_path.
+
+    Raises:
+        ImportError: If multiclean or rasterio is not installed.
+        FileNotFoundError: If input_path does not exist.
+
+    Example:
+        >>> from geoai.tools.multiclean import clean_raster
+        >>> clean_raster(
+        ...     "segmentation_raw.tif",
+        ...     "segmentation_cleaned.tif",
+        ...     class_values=[0, 1, 2],
+        ...     smooth_edge_size=3,
+        ...     min_island_size=50
+        ... )
+    """
+    check_multiclean_available()
+
+    if not RASTERIO_AVAILABLE:
+        raise ImportError(
+            "rasterio is required for raster operations. "
+            "Please install it with: pip install rasterio"
+        )
+
+    if not os.path.exists(input_path):
+        raise FileNotFoundError(f"Input file not found: {input_path}")
+
+    # Read input raster
+    with rasterio.open(input_path) as src:
+        # Read the specified band
+        mask = src.read(band)
+
+        # Get metadata
+        profile = src.profile.copy()
+
+        # Handle nodata
+        if nodata is None:
+            nodata = src.nodata
+
+        # Convert nodata to NaN if specified
+        if nodata is not None:
+            mask = mask.astype(np.float32)
+            mask[mask == nodata] = np.nan
+
+    # Clean the mask
+    cleaned = clean_segmentation_mask(
+        mask,
+        class_values=class_values,
+        smooth_edge_size=smooth_edge_size,
+        min_island_size=min_island_size,
+        connectivity=connectivity,
+        max_workers=max_workers,
+        fill_nan=fill_nan,
+    )
+
+    # Convert NaN back to nodata if needed
+    if nodata is not None:
+        # Convert any remaining NaN values back to nodata value
+        if np.isnan(cleaned).any():
+            cleaned = np.nan_to_num(cleaned, nan=nodata)
+
+    # Update profile for output
+    profile.update(
+        dtype=cleaned.dtype,
+        count=1,
+        compress="lzw",
+        nodata=nodata,
+    )
+
+    # Write cleaned raster
+    output_dir = os.path.dirname(os.path.abspath(output_path))
+    if output_dir and output_dir != os.path.abspath(os.sep):
+        os.makedirs(output_dir, exist_ok=True)
+    with rasterio.open(output_path, "w", **profile) as dst:
+        dst.write(cleaned, 1)
+
+
+def clean_raster_batch(
+    input_paths: List[str],
+    output_dir: str,
+    class_values: Optional[Union[int, List[int]]] = None,
+    smooth_edge_size: int = 2,
+    min_island_size: int = 100,
+    connectivity: int = 8,
+    max_workers: Optional[int] = None,
+    fill_nan: bool = False,
+    band: int = 1,
+    suffix: str = "_cleaned",
+    verbose: bool = True,
+) -> List[str]:
+    """
+    Clean multiple classification rasters in batch.
+
+    Processes multiple GeoTIFF files with the same cleaning parameters
+    and saves results to an output directory.
+
+    Args:
+        input_paths (list of str): List of paths to input GeoTIFF files.
+        output_dir (str): Directory to save cleaned files.
+        class_values (int, list of int, or None): Target class values.
+            Defaults to None (auto-detect).
+        smooth_edge_size (int): Kernel width for edge smoothing. Defaults to 2.
+        min_island_size (int): Minimum component area. Defaults to 100.
+        connectivity (int): Connectivity (4 or 8). Defaults to 8.
+        max_workers (int, optional): Thread pool size. Defaults to None.
+        fill_nan (bool): Whether to fill NaN pixels. Defaults to False.
+        band (int): Band index to read (1-indexed). Defaults to 1.
+        suffix (str): Suffix to add to output filenames. Defaults to "_cleaned".
+        verbose (bool): Whether to print progress. Defaults to True.
+
+    Returns:
+        list of str: Paths to cleaned output files.
+
+    Raises:
+        ImportError: If multiclean or rasterio is not installed.
+
+    Example:
+        >>> from geoai.tools.multiclean import clean_raster_batch
+        >>> input_files = ["mask1.tif", "mask2.tif", "mask3.tif"]
+        >>> outputs = clean_raster_batch(
+        ...     input_files,
+        ...     output_dir="cleaned_masks",
+        ...     min_island_size=50
+        ... )
+    """
+    check_multiclean_available()
+
+    # Create output directory
+    os.makedirs(output_dir, exist_ok=True)
+
+    output_paths = []
+
+    for i, input_path in enumerate(input_paths):
+        if verbose:
+            print(f"Processing {i+1}/{len(input_paths)}: {input_path}")
+
+        # Generate output filename
+        basename = os.path.basename(input_path)
+        name, ext = os.path.splitext(basename)
+        output_filename = f"{name}{suffix}{ext}"
+        output_path = os.path.join(output_dir, output_filename)
+
+        try:
+            # Clean the raster
+            clean_raster(
+                input_path,
+                output_path,
+                class_values=class_values,
+                smooth_edge_size=smooth_edge_size,
+                min_island_size=min_island_size,
+                connectivity=connectivity,
+                max_workers=max_workers,
+                fill_nan=fill_nan,
+                band=band,
+            )
+
+            output_paths.append(output_path)
+
+            if verbose:
+                print(f"  ✓ Saved to: {output_path}")
+
+        except Exception as e:
+            if verbose:
+                print(f"  ✗ Failed: {e}")
+            continue
+
+    return output_paths
+
+
+def compare_masks(
+    original: np.ndarray,
+    cleaned: np.ndarray,
+) -> Tuple[int, int, float]:
+    """
+    Compare original and cleaned masks to quantify changes.
+
+    Args:
+        original (np.ndarray): Original segmentation mask.
+        cleaned (np.ndarray): Cleaned segmentation mask.
+
+    Returns:
+        tuple: (pixels_changed, total_pixels, change_percentage)
+            - pixels_changed: Number of pixels that changed value
+            - total_pixels: Total number of valid pixels
+            - change_percentage: Percentage of pixels changed
+
+    Example:
+        >>> import numpy as np
+        >>> from geoai.tools.multiclean import compare_masks
+        >>> original = np.random.randint(0, 3, (512, 512))
+        >>> cleaned = original.copy()
+        >>> changed, total, pct = compare_masks(original, cleaned)
+        >>> print(f"Changed: {pct:.2f}%")
+    """
+    # Handle NaN values
+    valid_mask = ~(np.isnan(original) | np.isnan(cleaned))
+
+    # Count changed pixels
+    pixels_changed = np.sum((original != cleaned) & valid_mask)
+    total_pixels = np.sum(valid_mask)
+
+    # Calculate percentage
+    change_percentage = (pixels_changed / total_pixels * 100) if total_pixels > 0 else 0
+
+    return pixels_changed, total_pixels, change_percentage

geoai/train.py

@@ -1436,8 +1436,12 @@ def instance_segmentation_inference_on_geotiff(
         # Apply Non-Maximum Suppression to handle overlapping detections
         if len(all_detections) > 0:
             # Convert to tensors for NMS
-            boxes = torch.tensor([det["box"] for det in all_detections])
-            scores = torch.tensor([det["score"] for det in all_detections])
+            boxes = torch.tensor(
+                [det["box"] for det in all_detections], dtype=torch.float32
+            )
+            scores = torch.tensor(
+                [det["score"] for det in all_detections], dtype=torch.float32
+            )
 
             # Apply NMS with IoU threshold
             nms_threshold = 0.3  # IoU threshold for NMS
@@ -1917,6 +1921,96 @@ class SemanticRandomHorizontalFlip:
         return image, mask
 
 
+class SemanticRandomVerticalFlip:
+    """Random vertical flip transform for semantic segmentation."""
+
+    def __init__(self, prob: float = 0.5) -> None:
+        self.prob = prob
+
+    def __call__(
+        self, image: torch.Tensor, mask: torch.Tensor
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        if random.random() < self.prob:
+            # Flip image and mask along height dimension
+            image = torch.flip(image, dims=[1])
+            mask = torch.flip(mask, dims=[0])
+        return image, mask
+
+
+class SemanticRandomRotation90:
+    """Random 90-degree rotation transform for semantic segmentation."""
+
+    def __init__(self, prob: float = 0.5) -> None:
+        self.prob = prob
+
+    def __call__(
+        self, image: torch.Tensor, mask: torch.Tensor
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        if random.random() < self.prob:
+            # Randomly rotate by 90, 180, or 270 degrees
+            k = random.randint(1, 3)
+            image = torch.rot90(image, k, dims=[1, 2])
+            mask = torch.rot90(mask, k, dims=[0, 1])
+        return image, mask
+
+
+class SemanticBrightnessAdjustment:
+    """Random brightness adjustment transform for semantic segmentation."""
+
+    def __init__(
+        self, brightness_range: Tuple[float, float] = (0.8, 1.2), prob: float = 0.5
+    ) -> None:
+        """
+        Initialize brightness adjustment transform.
+
+        Args:
+            brightness_range: Tuple of (min, max) brightness factors.
+            prob: Probability of applying the transform.
+        """
+        self.brightness_range = brightness_range
+        self.prob = prob
+
+    def __call__(
+        self, image: torch.Tensor, mask: torch.Tensor
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        if random.random() < self.prob:
+            # Apply random brightness adjustment
+            factor = self.brightness_range[0] + random.random() * (
+                self.brightness_range[1] - self.brightness_range[0]
+            )
+            image = torch.clamp(image * factor, 0, 1)
+        return image, mask
+
+
+class SemanticContrastAdjustment:
+    """Random contrast adjustment transform for semantic segmentation."""
+
+    def __init__(
+        self, contrast_range: Tuple[float, float] = (0.8, 1.2), prob: float = 0.5
+    ) -> None:
+        """
+        Initialize contrast adjustment transform.
+
+        Args:
+            contrast_range: Tuple of (min, max) contrast factors.
+            prob: Probability of applying the transform.
+        """
+        self.contrast_range = contrast_range
+        self.prob = prob
+
+    def __call__(
+        self, image: torch.Tensor, mask: torch.Tensor
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        if random.random() < self.prob:
+            # Apply random contrast adjustment
+            factor = self.contrast_range[0] + random.random() * (
+                self.contrast_range[1] - self.contrast_range[0]
+            )
+            mean = image.mean(dim=(1, 2), keepdim=True)
+            image = torch.clamp((image - mean) * factor + mean, 0, 1)
+        return image, mask
+
+
 def get_semantic_transform(train: bool) -> Any:
     """
     Get transforms for semantic segmentation data augmentation.
@@ -2388,6 +2482,8 @@ def train_segmentation_model(
     resize_mode: str = "resize",
     num_workers: Optional[int] = None,
     early_stopping_patience: Optional[int] = None,
+    train_transforms: Optional[Callable] = None,
+    val_transforms: Optional[Callable] = None,
     **kwargs: Any,
 ) -> torch.nn.Module:
     """
@@ -2440,8 +2536,17 @@ def train_segmentation_model(
             'resize' - Resize images to target_size (may change aspect ratio)
             'pad' - Pad images to target_size (preserves aspect ratio). Defaults to 'resize'.
         num_workers (int): Number of workers for data loading. If None, uses 0 on macOS and Windows, 8 otherwise.
-        early_stopping_patience (int, optional): Number of epochs with no improvement after which
-            training will be stopped. If None, early stopping is disabled. Defaults to None.
+            Both image and mask should be torch.Tensor objects. The image tensor is expected to be in
+            CHW format (channels, height, width), and the mask tensor in HW format (height, width).
+            If None, uses default transforms (horizontal flip with 0.5 probability). Defaults to None.
+        val_transforms (callable, optional): Custom transforms for validation data.
+            Should be a callable that accepts (image, mask) tensors and returns transformed (image, mask).
+            The image tensor is expected to be in CHW format (channels, height, width), and the mask tensor in HW format (height, width).
+            Both image and mask should be torch.Tensor objects. If None, uses default transforms
+            (horizontal flip with 0.5 probability). Defaults to None.
+        val_transforms (callable, optional): Custom transforms for validation data.
+            Should be a callable that accepts (image, mask) tensors and returns transformed (image, mask).
+            If None, uses default transforms (no augmentation). Defaults to None.
         **kwargs: Additional arguments passed to smp.create_model().
     Returns:
         None: Model weights are saved to output_dir.
@@ -2584,10 +2689,22 @@ def train_segmentation_model(
                 print("No resizing needed.")
 
     # Create datasets
+    # Use custom transforms if provided, otherwise use default transforms
+    train_transform = (
+        train_transforms
+        if train_transforms is not None
+        else get_semantic_transform(train=True)
+    )
+    val_transform = (
+        val_transforms
+        if val_transforms is not None
+        else get_semantic_transform(train=False)
+    )
+
     train_dataset = SemanticSegmentationDataset(
         train_imgs,
         train_labels,
-        transforms=get_semantic_transform(train=True),
+        transforms=train_transform,
         num_channels=num_channels,
         target_size=target_size,
         resize_mode=resize_mode,
@@ -2596,7 +2713,7 @@ def train_segmentation_model(
     val_dataset = SemanticSegmentationDataset(
         val_imgs,
         val_labels,
-        transforms=get_semantic_transform(train=False),
+        transforms=val_transform,
         num_channels=num_channels,
         target_size=target_size,
         resize_mode=resize_mode,