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

geoai/__init__.py

@@ -2,7 +2,7 @@
 
 __author__ = """Qiusheng Wu"""
 __email__ = "giswqs@gmail.com"
-__version__ = "0.17.0"
+__version__ = "0.18.0"
 
 
 import os
@@ -121,3 +121,18 @@ from .timm_segment import (
     timm_semantic_segmentation,
     push_timm_model_to_hub,
 )
+
+# Import tools subpackage
+from . import tools
+
+# Expose commonly used tools at package level for convenience
+try:
+    from .tools import (
+        clean_segmentation_mask,
+        clean_raster,
+        clean_raster_batch,
+        compare_masks,
+    )
+except ImportError:
+    # MultiClean not available (missing dependency)
+    pass

geoai/agents/geo_agents.py

@@ -55,7 +55,9 @@ class UICallbackHandler:
                 # Make tool names more user-friendly
                 friendly_name = tool_name.replace("_", " ").title()
                 self.status_widget.value = (
-                    f"<span style='color:#0a7'>⚙️ {friendly_name}...</span>"
+                    f"<span style='color:#0a7'>"
+                    f"<i class='fas fa-spinner fa-spin' style='font-size:1.2em'></i> "
+                    f"{friendly_name}...</span>"
                 )
 
 
@@ -396,7 +398,10 @@ class GeoAgent(Agent):
         btn_clear = widgets.Button(
             description="Clear", icon="trash", layout=widgets.Layout(width="120px")
         )
-        status = widgets.HTML("<span style='color:#666'>Ready.</span>")
+        status = widgets.HTML(
+            "<link rel='stylesheet' href='https://cdnjs.cloudflare.com/ajax/libs/font-awesome/6.4.0/css/all.min.css'>"
+            "<span style='color:#666'>Ready.</span>"
+        )
 
         examples = widgets.Dropdown(
             options=[
@@ -1017,7 +1022,10 @@ CRITICAL: Return ONLY JSON. NO explanatory text, NO made-up data."""
         btn_clear = widgets.Button(
             description="Clear", icon="trash", layout=widgets.Layout(width="120px")
         )
-        status = widgets.HTML("<span style='color:#666'>Ready to search.</span>")
+        status = widgets.HTML(
+            "<link rel='stylesheet' href='https://cdnjs.cloudflare.com/ajax/libs/font-awesome/6.4.0/css/all.min.css'>"
+            "<span style='color:#666'>Ready to search.</span>"
+        )
 
         examples = widgets.Dropdown(
             options=[

geoai/change_detection.py

@@ -13,7 +13,8 @@ from skimage.transform import resize
 try:
     from torchange.models.segment_any_change import AnyChange, show_change_masks
 except ImportError:
-    print("torchange requires Python 3.11 or higher")
+    AnyChange = None
+    show_change_masks = None
 
 from .utils import download_file
 
@@ -36,6 +37,13 @@ class ChangeDetection:
 
     def _init_model(self):
         """Initialize the AnyChange model."""
+        if AnyChange is None:
+            raise ImportError(
+                "The 'torchange' package is required for change detection. "
+                "Please install it using: pip install torchange\n"
+                "Note: torchange requires Python 3.11 or higher."
+            )
+
         if self.sam_checkpoint is None:
             self.sam_checkpoint = download_checkpoint(self.sam_model_type)
 
@@ -551,6 +559,13 @@ class ChangeDetection:
         Returns:
             matplotlib.figure.Figure: The figure object
         """
+        if show_change_masks is None:
+            raise ImportError(
+                "The 'torchange' package is required for change detection visualization. "
+                "Please install it using: pip install torchange\n"
+                "Note: torchange requires Python 3.11 or higher."
+            )
+
         change_masks, img1, img2 = self.detect_changes(
             image1_path, image2_path, return_results=True
         )

geoai/timm_segment.py

@@ -241,7 +241,10 @@ class TimmSegmentationModel(pl.LightningModule):
         )
 
         scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
-            optimizer, mode="min", factor=0.5, patience=5, verbose=True
+            optimizer,
+            mode="min",
+            factor=0.5,
+            patience=5,
         )
 
         return {

geoai/tools/__init__.py

@@ -0,0 +1,65 @@
+"""
+GeoAI Tools - Utility functions and integrations for geospatial AI workflows.
+
+This subpackage contains various tools and integrations for enhancing
+geospatial AI workflows, including post-processing utilities and
+third-party library integrations.
+"""
+
+__all__ = []
+
+# MultiClean integration (optional dependency)
+try:
+    from .multiclean import (
+        clean_segmentation_mask,
+        clean_raster,
+        clean_raster_batch,
+        compare_masks,
+        check_multiclean_available,
+    )
+
+    __all__.extend(
+        [
+            "clean_segmentation_mask",
+            "clean_raster",
+            "clean_raster_batch",
+            "compare_masks",
+            "check_multiclean_available",
+        ]
+    )
+except ImportError:
+    # MultiClean not installed - functions will not be available
+    pass
+
+# OmniCloudMask integration (optional dependency)
+try:
+    from .cloudmask import (
+        predict_cloud_mask,
+        predict_cloud_mask_from_raster,
+        predict_cloud_mask_batch,
+        calculate_cloud_statistics,
+        create_cloud_free_mask,
+        check_omnicloudmask_available,
+        CLEAR,
+        THICK_CLOUD,
+        THIN_CLOUD,
+        CLOUD_SHADOW,
+    )
+
+    __all__.extend(
+        [
+            "predict_cloud_mask",
+            "predict_cloud_mask_from_raster",
+            "predict_cloud_mask_batch",
+            "calculate_cloud_statistics",
+            "create_cloud_free_mask",
+            "check_omnicloudmask_available",
+            "CLEAR",
+            "THICK_CLOUD",
+            "THIN_CLOUD",
+            "CLOUD_SHADOW",
+        ]
+    )
+except ImportError:
+    # OmniCloudMask not installed - functions will not be available
+    pass

geoai/tools/cloudmask.py

@@ -0,0 +1,431 @@
+"""
+OmniCloudMask integration for cloud and cloud shadow detection in satellite imagery.
+
+This module provides functions to use OmniCloudMask (https://github.com/DPIRD-DMA/OmniCloudMask)
+for detecting clouds and cloud shadows in satellite imagery. OmniCloudMask performs semantic
+segmentation to classify pixels into: Clear (0), Thick Cloud (1), Thin Cloud (2), Cloud Shadow (3).
+
+Supports Sentinel-2, Landsat 8, PlanetScope, and Maxar imagery at 10-50m resolution.
+"""
+
+import os
+from typing import Optional, List, Tuple, Dict, Any
+import numpy as np
+
+try:
+    from omnicloudmask import predict_from_array
+
+    OMNICLOUDMASK_AVAILABLE = True
+except ImportError:
+    OMNICLOUDMASK_AVAILABLE = False
+
+try:
+    import rasterio
+
+    RASTERIO_AVAILABLE = True
+except ImportError:
+    RASTERIO_AVAILABLE = False
+
+
+# Cloud mask class values
+CLEAR = 0
+THICK_CLOUD = 1
+THIN_CLOUD = 2
+CLOUD_SHADOW = 3
+
+
+def check_omnicloudmask_available():
+    """
+    Check if omnicloudmask is installed.
+
+    Raises:
+        ImportError: If omnicloudmask is not installed.
+    """
+    if not OMNICLOUDMASK_AVAILABLE:
+        raise ImportError(
+            "omnicloudmask is not installed. "
+            "Please install it with: pip install omnicloudmask "
+            "or: pip install geoai-py[extra]"
+        )
+
+
+def predict_cloud_mask(
+    image: np.ndarray,
+    batch_size: int = 1,
+    inference_device: str = "cpu",
+    inference_dtype: str = "fp32",
+    patch_size: int = 1000,
+    export_confidence: bool = False,
+    model_version: int = 3,
+) -> np.ndarray:
+    """
+    Predict cloud mask from a numpy array using OmniCloudMask.
+
+    This function classifies each pixel into one of four categories:
+    - 0: Clear
+    - 1: Thick Cloud
+    - 2: Thin Cloud
+    - 3: Cloud Shadow
+
+    Args:
+        image (np.ndarray): Input image array with shape (3, height, width) or (height, width, 3).
+            Should contain Red, Green, and NIR bands. Values should be in reflectance (0-1)
+            or digital numbers (0-10000 typical for Sentinel-2/Landsat).
+        batch_size (int): Number of patches to process per inference batch. Defaults to 1.
+        inference_device (str): Device for inference ('cpu', 'cuda', or 'mps'). Defaults to 'cpu'.
+        inference_dtype (str): Data type for inference ('fp32', 'fp16', or 'bf16').
+            'bf16' recommended for speed on compatible hardware. Defaults to 'fp32'.
+        patch_size (int): Size of patches for processing large images. Defaults to 1000.
+        export_confidence (bool): If True, also returns confidence map. Defaults to False.
+        model_version (int): Model version to use (1, 2, or 3). Defaults to 3.
+
+    Returns:
+        np.ndarray: Cloud mask array with shape (height, width) containing class predictions.
+            If export_confidence=True, returns tuple of (mask, confidence).
+
+    Raises:
+        ImportError: If omnicloudmask is not installed.
+        ValueError: If image has wrong shape or number of channels.
+
+    Example:
+        >>> import numpy as np
+        >>> from geoai.tools.cloudmask import predict_cloud_mask
+        >>> # Create synthetic image (3 bands: R, G, NIR)
+        >>> image = np.random.rand(3, 512, 512) * 10000
+        >>> mask = predict_cloud_mask(image)
+        >>> print(f"Clear pixels: {(mask == 0).sum()}")
+    """
+    check_omnicloudmask_available()
+
+    # Ensure image has correct shape (3, H, W)
+    if image.ndim != 3:
+        raise ValueError(f"Image must be 3D, got shape {image.shape}")
+
+    # Convert (H, W, 3) to (3, H, W) if needed
+    if image.shape[2] == 3 and image.shape[0] != 3:
+        image = np.transpose(image, (2, 0, 1))
+
+    if image.shape[0] != 3:
+        raise ValueError(
+            f"Image must have 3 channels (R, G, NIR), got {image.shape[0]} channels"
+        )
+
+    # Call OmniCloudMask
+    result = predict_from_array(
+        image,
+        batch_size=batch_size,
+        inference_device=inference_device,
+        inference_dtype=inference_dtype,
+        patch_size=patch_size,
+        export_confidence=export_confidence,
+        model_version=model_version,
+    )
+
+    # Handle output shape - omnicloudmask returns (1, H, W) or ((1, H, W), (1, H, W))
+    if export_confidence:
+        mask, confidence = result
+        # Squeeze batch dimension
+        mask = mask.squeeze(0) if mask.ndim == 3 else mask
+        confidence = confidence.squeeze(0) if confidence.ndim == 3 else confidence
+        return mask, confidence
+    else:
+        # Squeeze batch dimension
+        return result.squeeze(0) if result.ndim == 3 else result
+
+
+def predict_cloud_mask_from_raster(
+    input_path: str,
+    output_path: str,
+    red_band: int = 1,
+    green_band: int = 2,
+    nir_band: int = 3,
+    batch_size: int = 1,
+    inference_device: str = "cpu",
+    inference_dtype: str = "fp32",
+    patch_size: int = 1000,
+    export_confidence: bool = False,
+    model_version: int = 3,
+) -> None:
+    """
+    Predict cloud mask from a GeoTIFF file and save the result.
+
+    Reads a multi-band raster, extracts RGB+NIR bands, applies OmniCloudMask,
+    and saves the result while preserving geospatial metadata.
+
+    Args:
+        input_path (str): Path to input GeoTIFF file.
+        output_path (str): Path to save cloud mask GeoTIFF.
+        red_band (int): Band index for Red (1-indexed). Defaults to 1.
+        green_band (int): Band index for Green (1-indexed). Defaults to 2.
+        nir_band (int): Band index for NIR (1-indexed). Defaults to 3.
+        batch_size (int): Patches per inference batch. Defaults to 1.
+        inference_device (str): Device ('cpu', 'cuda', 'mps'). Defaults to 'cpu'.
+        inference_dtype (str): Dtype ('fp32', 'fp16', 'bf16'). Defaults to 'fp32'.
+        patch_size (int): Patch size for large images. Defaults to 1000.
+        export_confidence (bool): Export confidence map. Defaults to False.
+        model_version (str): Model version ('1.0', '2.0', '3.0'). Defaults to '3.0'.
+
+    Returns:
+        None: Writes cloud mask to output_path.
+
+    Raises:
+        ImportError: If omnicloudmask or rasterio not installed.
+        FileNotFoundError: If input_path doesn't exist.
+
+    Example:
+        >>> from geoai.tools.cloudmask import predict_cloud_mask_from_raster
+        >>> predict_cloud_mask_from_raster(
+        ...     "sentinel2_image.tif",
+        ...     "cloud_mask.tif",
+        ...     red_band=4,   # Sentinel-2 band order
+        ...     green_band=3,
+        ...     nir_band=8
+        ... )
+    """
+    check_omnicloudmask_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 required bands
+        red = src.read(red_band).astype(np.float32)
+        green = src.read(green_band).astype(np.float32)
+        nir = src.read(nir_band).astype(np.float32)
+
+        # Stack into (3, H, W)
+        image = np.stack([red, green, nir], axis=0)
+
+        # Get metadata
+        profile = src.profile.copy()
+
+    # Predict cloud mask
+    result = predict_cloud_mask(
+        image,
+        batch_size=batch_size,
+        inference_device=inference_device,
+        inference_dtype=inference_dtype,
+        patch_size=patch_size,
+        export_confidence=export_confidence,
+        model_version=model_version,
+    )
+
+    # Handle confidence output
+    if export_confidence:
+        mask, confidence = result
+    else:
+        mask = result
+
+    # Update profile for output
+    profile.update(
+        dtype=np.uint8,
+        count=1,
+        compress="lzw",
+        nodata=None,
+    )
+
+    # Write cloud mask
+    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(mask.astype(np.uint8), 1)
+
+    # Optionally write confidence map
+    if export_confidence:
+        confidence_path = output_path.replace(".tif", "_confidence.tif")
+        profile.update(dtype=np.float32)
+        with rasterio.open(confidence_path, "w", **profile) as dst:
+            dst.write(confidence, 1)
+
+
+def predict_cloud_mask_batch(
+    input_paths: List[str],
+    output_dir: str,
+    red_band: int = 1,
+    green_band: int = 2,
+    nir_band: int = 3,
+    batch_size: int = 1,
+    inference_device: str = "cpu",
+    inference_dtype: str = "fp32",
+    patch_size: int = 1000,
+    export_confidence: bool = False,
+    model_version: int = 3,
+    suffix: str = "_cloudmask",
+    verbose: bool = True,
+) -> List[str]:
+    """
+    Predict cloud masks for multiple rasters in batch.
+
+    Processes multiple GeoTIFF files with the same cloud detection parameters
+    and saves results to an output directory.
+
+    Args:
+        input_paths (list of str): Paths to input GeoTIFF files.
+        output_dir (str): Directory to save cloud masks.
+        red_band (int): Red band index. Defaults to 1.
+        green_band (int): Green band index. Defaults to 2.
+        nir_band (int): NIR band index. Defaults to 3.
+        batch_size (int): Patches per batch. Defaults to 1.
+        inference_device (str): Device. Defaults to 'cpu'.
+        inference_dtype (str): Dtype. Defaults to 'fp32'.
+        patch_size (int): Patch size. Defaults to 1000.
+        export_confidence (bool): Export confidence. Defaults to False.
+        model_version (str): Model version. Defaults to '3.0'.
+        suffix (str): Suffix for output filenames. Defaults to '_cloudmask'.
+        verbose (bool): Print progress. Defaults to True.
+
+    Returns:
+        list of str: Paths to output cloud mask files.
+
+    Raises:
+        ImportError: If omnicloudmask or rasterio not installed.
+
+    Example:
+        >>> from geoai.tools.cloudmask import predict_cloud_mask_batch
+        >>> files = ["scene1.tif", "scene2.tif", "scene3.tif"]
+        >>> outputs = predict_cloud_mask_batch(
+        ...     files,
+        ...     output_dir="cloud_masks",
+        ...     inference_device="cuda"
+        ... )
+    """
+    check_omnicloudmask_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:
+            # Predict cloud mask
+            predict_cloud_mask_from_raster(
+                input_path,
+                output_path,
+                red_band=red_band,
+                green_band=green_band,
+                nir_band=nir_band,
+                batch_size=batch_size,
+                inference_device=inference_device,
+                inference_dtype=inference_dtype,
+                patch_size=patch_size,
+                export_confidence=export_confidence,
+                model_version=model_version,
+            )
+
+            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 calculate_cloud_statistics(
+    mask: np.ndarray,
+) -> Dict[str, Any]:
+    """
+    Calculate statistics from a cloud mask.
+
+    Args:
+        mask (np.ndarray): Cloud mask array with values 0-3.
+
+    Returns:
+        dict: Statistics including:
+            - total_pixels: Total number of pixels
+            - clear_pixels: Number of clear pixels
+            - thick_cloud_pixels: Number of thick cloud pixels
+            - thin_cloud_pixels: Number of thin cloud pixels
+            - shadow_pixels: Number of cloud shadow pixels
+            - clear_percent: Percentage of clear pixels
+            - cloud_percent: Percentage of cloudy pixels (thick + thin)
+            - shadow_percent: Percentage of shadow pixels
+
+    Example:
+        >>> from geoai.tools.cloudmask import calculate_cloud_statistics
+        >>> import numpy as np
+        >>> mask = np.random.randint(0, 4, (512, 512))
+        >>> stats = calculate_cloud_statistics(mask)
+        >>> print(f"Clear: {stats['clear_percent']:.1f}%")
+    """
+    total_pixels = mask.size
+
+    clear_pixels = (mask == CLEAR).sum()
+    thick_cloud_pixels = (mask == THICK_CLOUD).sum()
+    thin_cloud_pixels = (mask == THIN_CLOUD).sum()
+    shadow_pixels = (mask == CLOUD_SHADOW).sum()
+
+    cloud_pixels = thick_cloud_pixels + thin_cloud_pixels
+
+    return {
+        "total_pixels": int(total_pixels),
+        "clear_pixels": int(clear_pixels),
+        "thick_cloud_pixels": int(thick_cloud_pixels),
+        "thin_cloud_pixels": int(thin_cloud_pixels),
+        "shadow_pixels": int(shadow_pixels),
+        "clear_percent": float(clear_pixels / total_pixels * 100),
+        "cloud_percent": float(cloud_pixels / total_pixels * 100),
+        "shadow_percent": float(shadow_pixels / total_pixels * 100),
+    }
+
+
+def create_cloud_free_mask(
+    mask: np.ndarray,
+    include_thin_clouds: bool = False,
+    include_shadows: bool = False,
+) -> np.ndarray:
+    """
+    Create a binary mask of cloud-free pixels.
+
+    Args:
+        mask (np.ndarray): Cloud mask with values 0-3.
+        include_thin_clouds (bool): If True, treats thin clouds as acceptable.
+            Defaults to False.
+        include_shadows (bool): If True, treats shadows as acceptable.
+            Defaults to False.
+
+    Returns:
+        np.ndarray: Binary mask where 1 = usable, 0 = not usable.
+
+    Example:
+        >>> from geoai.tools.cloudmask import create_cloud_free_mask
+        >>> import numpy as np
+        >>> mask = np.random.randint(0, 4, (512, 512))
+        >>> cloud_free = create_cloud_free_mask(mask)
+        >>> print(f"Usable pixels: {cloud_free.sum()}")
+    """
+    # Start with clear pixels
+    usable = mask == CLEAR
+
+    # Optionally include thin clouds
+    if include_thin_clouds:
+        usable = usable | (mask == THIN_CLOUD)
+
+    # Optionally include shadows
+    if include_shadows:
+        usable = usable | (mask == CLOUD_SHADOW)
+
+    return usable.astype(np.uint8)

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,

{geoai_py-0.17.0.dist-info → geoai_py-0.18.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: geoai-py
-Version: 0.17.0
+Version: 0.18.0
 Summary: A Python package for using Artificial Intelligence (AI) with geospatial data
 Author-email: Qiusheng Wu <giswqs@gmail.com>
 License: MIT License
@@ -39,15 +39,18 @@ Requires-Dist: rioxarray
 Requires-Dist: scikit-image
 Requires-Dist: scikit-learn
 Requires-Dist: timm
+Requires-Dist: tokenizers>=0.22.1
 Requires-Dist: torch
 Requires-Dist: torchgeo
 Requires-Dist: torchinfo
 Requires-Dist: tqdm
-Requires-Dist: transformers
+Requires-Dist: transformers>=4.57.1
 Provides-Extra: extra
 Requires-Dist: overturemaps; extra == "extra"
 Requires-Dist: torchange; extra == "extra"
 Requires-Dist: lightly-train; extra == "extra"
+Requires-Dist: multiclean; extra == "extra"
+Requires-Dist: omnicloudmask; extra == "extra"
 Provides-Extra: agents
 Requires-Dist: strands-agents; extra == "agents"
 Requires-Dist: strands-agents-tools; extra == "agents"

{geoai_py-0.17.0.dist-info → geoai_py-0.18.0.dist-info}/RECORD

@@ -1,5 +1,5 @@
-geoai/__init__.py,sha256=PFtD_WPr8ZKYK7z_TtN8ZRvdCeFD4Te-eXzuqD5ma8o,4287
-geoai/change_detection.py,sha256=B0eawgiQYxEC51lI1q6V83K1mRaAVW8-Kf2GaVhOTBQ,59799
+geoai/__init__.py,sha256=9wW42cjIMMY_ec5s-DOO04XgQHXVre8CisL72wvJCIU,4620
+geoai/change_detection.py,sha256=pdQofnPRiwoES8vMln2vHghRnpeTdsmqLir74dnqZYU,60389
 geoai/classify.py,sha256=0DcComVR6vKU4qWtH2oHVeXc7ZTcV0mFvdXRtlNmolo,35637
 geoai/detectron2.py,sha256=dOOFM9M9-6PV8q2A4-mnIPrz7yTo-MpEvDiAW34nl0w,14610
 geoai/dinov3.py,sha256=u4Lulihhvs4wTgi84RjRw8jWQpB8omQSl-dVVryNVus,40377
@@ -11,20 +11,23 @@ geoai/map_widgets.py,sha256=QLmkILsztNaRXRULHKOd7Glb7S0pEWXSK9-P8S5AuzQ,5856
 geoai/sam.py,sha256=O6S-kGiFn7YEcFbfWFItZZQOhnsm6-GlunxQLY0daEs,34345
 geoai/segment.py,sha256=yBGTxA-ti8lBpk7WVaBOp6yP23HkaulKJQk88acrmZ0,43788
 geoai/segmentation.py,sha256=7yEzBSKCyHW1dNssoK0rdvhxi2IXsIQIFSga817KdI4,11535
-geoai/timm_segment.py,sha256=G37UmSKKR_6CJPkt0TTj1Ik9_xd5i2fQ79W_DA3dheM,38447
+geoai/timm_segment.py,sha256=GfvWmxT6t1S99-iZOf8PlsCkwodIUyrt0AwO_j6dCjE,38470
 geoai/timm_train.py,sha256=y_Sm9Fwe7bTsHEKdtPee5rGY7s01CbkAZKP1TwUDXlU,20551
-geoai/train.py,sha256=aWMIen_OPSQ5UwwgURVa3ycqtdZfzuZ0ign2_OHnia8,170310
+geoai/train.py,sha256=Ef-lCCQvaMWl3wvhi-IYYi9sdR4YBqMt9QkfiRAUlkQ,174762
 geoai/utils.py,sha256=AUdVj1tt864UFxJtsatpUmXRV9-Lw4f4tbdyjqj0c3c,360240
 geoai/agents/__init__.py,sha256=5xtb_dGpI26nPFcAm8Dj7O4bLskqr1xTw2BRQqbgH4w,285
 geoai/agents/catalog_models.py,sha256=19E-PiE7FvpGEiOi4gDMKPf257FOhLseuVGWJbOjrDs,2089
 geoai/agents/catalog_tools.py,sha256=psVw7-di65hhnJUFqWXFoOkbGaG2_sHrQhA5vdXp3x4,33597
-geoai/agents/geo_agents.py,sha256=sZSVTO1A3X_iz1D4W7SnAPgYhvIhpp0g882y5rr3P0w,58916
+geoai/agents/geo_agents.py,sha256=2O25JTNN2qr9YXhjHJfXFMByA7B3yyhGGzWyzroIDYg,59302
 geoai/agents/map_tools.py,sha256=OK5uB0VUHjjUnc-DYRy2CQ__kyUIARSCPBucGabO0Xw,60669
 geoai/agents/stac_models.py,sha256=N2kv7HHdAKT8lgYheEd98QK4l0UcbgpNTLOWW24ayBs,2573
 geoai/agents/stac_tools.py,sha256=ILUg2xFRXVZ9WHOfPeJBvPSFT7lRsPLnGMZhnpDZ1co,16107
-geoai_py-0.17.0.dist-info/licenses/LICENSE,sha256=TlBm8mRusRVB9yF2NTg-STcb71v69-XZaKaPdshqP2I,1074
-geoai_py-0.17.0.dist-info/METADATA,sha256=IvRWtKnxTxFlsdUTs0QJfC_WsrHi3Cn3zOhgb3Qw-pY,11122
-geoai_py-0.17.0.dist-info/WHEEL,sha256=JNWh1Fm1UdwIQV075glCn4MVuCRs0sotJIq-J6rbxCU,109
-geoai_py-0.17.0.dist-info/entry_points.txt,sha256=uGp3Az3HURIsRHP9v-ys0hIbUuBBNUfXv6VbYHIXeg4,41
-geoai_py-0.17.0.dist-info/top_level.txt,sha256=1YkCUWu-ii-0qIex7kbwAvfei-gos9ycyDyUCJPNWHY,6
-geoai_py-0.17.0.dist-info/RECORD,,
+geoai/tools/__init__.py,sha256=McC49tQjxrTha1TS69IeM3rRvqVQP3H1NdAZPZPpKEI,1683
+geoai/tools/cloudmask.py,sha256=qzvqVa8FAEgd8mePXBaV5Ptx4fHhwfS1BsYL0JAZBjM,14500
+geoai/tools/multiclean.py,sha256=TVwmWgeQyGIyUuCe10b6pGCtgIl8TkZmcgVXPimn9uM,11949
+geoai_py-0.18.0.dist-info/licenses/LICENSE,sha256=TlBm8mRusRVB9yF2NTg-STcb71v69-XZaKaPdshqP2I,1074
+geoai_py-0.18.0.dist-info/METADATA,sha256=iqGhGU99OgqXBXfMF5r4GfyrByGQFFzx3PHBOAxxYG0,11255
+geoai_py-0.18.0.dist-info/WHEEL,sha256=JNWh1Fm1UdwIQV075glCn4MVuCRs0sotJIq-J6rbxCU,109
+geoai_py-0.18.0.dist-info/entry_points.txt,sha256=uGp3Az3HURIsRHP9v-ys0hIbUuBBNUfXv6VbYHIXeg4,41
+geoai_py-0.18.0.dist-info/top_level.txt,sha256=1YkCUWu-ii-0qIex7kbwAvfei-gos9ycyDyUCJPNWHY,6
+geoai_py-0.18.0.dist-info/RECORD,,