Upload 3 files

load.py

@@ -0,0 +1,358 @@
+"""
+Load and inference functions for MSS Cloud Detection Model
+Compatible with mlstac package
+"""
+
+import torch
+import torch.nn as nn
+import numpy as np
+from pathlib import Path
+from typing import Tuple, Optional
+import pytorch_lightning as pl
+import segmentation_models_pytorch as smp
+from tqdm import tqdm
+
+
+# ============================================================================
+# MODEL DEFINITION (copied from your model.py)
+# ============================================================================
+
+class MSSSegmentationModel(pl.LightningModule):
+    """UNet para cloud segmentation en MSS."""
+    
+    def __init__(
+        self,
+        in_channels: int = 4,
+        num_classes: int = 4,
+        encoder: str = "efficientnet-b3",
+        lr: float = 3e-4,
+        weight_decay: float = 1e-4,
+    ):
+        super().__init__()
+        self.save_hyperparameters()
+        
+        self.model = smp.Unet(
+            encoder_name=encoder,
+            encoder_weights=None,
+            in_channels=in_channels,
+            classes=num_classes,
+            encoder_depth=5,
+            activation=None,
+            decoder_attention_type="scse",
+        )
+
+    def forward(self, x):
+        return self.model(x)
+
+
+# ============================================================================
+# INFERENCE UTILITIES
+# ============================================================================
+
+def get_spline_window(size: int, power: int = 2) -> np.ndarray:
+    """Generate Hann window for smooth blending."""
+    intersection = np.hanning(size)
+    window_2d = np.outer(intersection, intersection)
+    return (window_2d ** power).astype(np.float32)
+
+
+def apply_physical_rules(
+    pred: np.ndarray,
+    image: np.ndarray,
+    merge_clouds: bool = False,
+    saturation_threshold: float = 0.35,
+) -> np.ndarray:
+    """
+    Apply physical rules for better cloud detection.
+    
+    Args:
+        pred: Predicted classes (H, W)
+        image: Input image (4, H, W) in reflectance [0, 1]
+        merge_clouds: If True, merge thin+thick into single cloud class
+        saturation_threshold: Threshold for detecting saturated bright clouds
+    """
+    pred = pred.copy()
+    
+    # Mask nodata pixels
+    nodata_mask = np.all(image == 0, axis=0)
+    pred[nodata_mask] = 0
+    
+    # Detect very bright pixels (likely thick clouds)
+    bright_b0 = image[0] > saturation_threshold
+    bright_b1 = image[1] > saturation_threshold * 0.80
+    saturated_mask = bright_b0 & bright_b1
+    
+    if merge_clouds:
+        # Set to cloud (1)
+        pred[saturated_mask] = 1
+    else:
+        # Set to thick cloud (2)
+        pred[saturated_mask] = 2
+    
+    return pred
+
+
+# ============================================================================
+# MLSTAC-COMPATIBLE FUNCTIONS
+# ============================================================================
+
+def compiled_model(
+    model_dir: Path,
+    stac_item=None,
+    device: str = "cpu",
+    merge_clouds: bool = False,
+    **kwargs
+) -> nn.Module:
+    """
+    Load compiled model for inference.
+    
+    Args:
+        model_dir: Directory containing the .ckpt file
+        stac_item: STAC item metadata (optional)
+        device: 'cpu' or 'cuda'
+        merge_clouds: If True, output will have 3 classes (clear, cloud, shadow)
+                     If False, output will have 4 classes (clear, thin, thick, shadow)
+    
+    Returns:
+        Loaded model in eval mode
+    """
+    # Find checkpoint file
+    ckpt_files = list(model_dir.glob("*.ckpt"))
+    if not ckpt_files:
+        raise FileNotFoundError(f"No .ckpt file found in {model_dir}")
+    
+    ckpt_path = ckpt_files[0]
+    
+    # Load model
+    model = MSSSegmentationModel.load_from_checkpoint(
+        ckpt_path,
+        map_location=device
+    )
+    model.eval()
+    model.to(device)
+    
+    # Disable gradients
+    for param in model.parameters():
+        param.requires_grad = False
+    
+    # Store merge_clouds flag for predict_large
+    model.merge_clouds = merge_clouds
+    
+    print(f"✅ Model loaded from {ckpt_path.name}")
+    print(f"   Device: {device}")
+    print(f"   Classes: {'3 (merged)' if merge_clouds else '4 (original)'}")
+    
+    return model
+
+
+def predict_large(
+    image: np.ndarray,
+    model: nn.Module,
+    chunk_size: int = 512,
+    overlap: int = 256,
+    batch_size: int = 1,
+    device: str = "cpu",
+    nodata: float = 0.0,
+    apply_rules: bool = True,
+    saturation_threshold: float = 0.35,
+    **kwargs
+) -> np.ndarray:
+    """
+    Predict on large images using sliding window with overlap blending.
+    
+    Args:
+        image: Input image (C, H, W) in reflectance [0, 1]
+        model: Loaded model from compiled_model()
+        chunk_size: Size of inference tiles (default: 1024)
+        overlap: Overlap between tiles for smooth blending (default: 256)
+        batch_size: Number of tiles to process in parallel (default: 1)
+        device: 'cpu' or 'cuda'
+        nodata: Value representing no-data pixels
+        apply_rules: Whether to apply physical rules post-processing
+        saturation_threshold: Threshold for detecting bright clouds
+    
+    Returns:
+        Predicted class labels (H, W) with shape matching input
+        - If merge_clouds=False: 0=clear, 1=thin, 2=thick, 3=shadow
+        - If merge_clouds=True: 0=clear, 1=cloud, 2=shadow
+    """
+    model.eval()
+    model.to(device)
+    
+    merge_clouds = getattr(model, 'merge_clouds', False)
+    
+    C, H, W = image.shape
+    
+    # Direct inference for small images
+    if H <= chunk_size and W <= chunk_size:
+        with torch.no_grad():
+            img_tensor = torch.from_numpy(image).unsqueeze(0).float().to(device)
+            logits = model(img_tensor)
+            
+            if merge_clouds:
+                # Merge thin(1) + thick(2) probabilities
+                probs = torch.softmax(logits, dim=1)
+                probs_merged = torch.zeros(1, 3, H, W, device=device)
+                probs_merged[:, 0] = probs[:, 0]  # clear
+                probs_merged[:, 1] = probs[:, 1] + probs[:, 2]  # cloud
+                probs_merged[:, 2] = probs[:, 3]  # shadow
+                pred = probs_merged.argmax(1).squeeze().cpu().numpy().astype(np.uint8)
+            else:
+                pred = logits.argmax(1).squeeze().cpu().numpy().astype(np.uint8)
+        
+        if apply_rules:
+            pred = apply_physical_rules(pred, image, merge_clouds, saturation_threshold)
+        
+        return pred
+    
+    # Sliding window for large images
+    step = chunk_size - overlap
+    half_tile = chunk_size // 2
+    
+    # Pad image
+    image_padded = np.pad(
+        image,
+        ((0, 0), (half_tile, half_tile + chunk_size), (half_tile, half_tile + chunk_size)),
+        mode="reflect"
+    )
+    
+    _, H_pad, W_pad = image_padded.shape
+    
+    # Initialize accumulators - ALWAYS 4 classes, merge at the end if needed
+    num_classes = 4
+    probs_sum = np.zeros((num_classes, H_pad, W_pad), dtype=np.float32)
+    weight_sum = np.zeros((H_pad, W_pad), dtype=np.float32)
+    
+    # Blending window
+    window = get_spline_window(chunk_size, power=2)
+    
+    # Generate tile coordinates
+    coords = [
+        (r, c)
+        for r in range(0, H_pad - chunk_size + 1, step)
+        for c in range(0, W_pad - chunk_size + 1, step)
+    ]
+    
+    # Process tiles in batches
+    with torch.no_grad():
+        for i in tqdm(range(0, len(coords), batch_size), desc="  Tiles", leave=False, disable=True):
+            batch_coords = coords[i:i + batch_size]
+            
+            # Extract tiles
+            tiles = np.stack([
+                image_padded[:, r:r + chunk_size, c:c + chunk_size]
+                for r, c in batch_coords
+            ])
+            
+            # Inference
+            tiles_tensor = torch.from_numpy(tiles).float().to(device)
+            logits = model(tiles_tensor)
+            probs = torch.softmax(logits, dim=1).cpu().numpy()
+            
+            # Accumulate with blending - ALWAYS accumulate 4 classes
+            for j, (r, c) in enumerate(batch_coords):
+                probs_sum[:, r:r + chunk_size, c:c + chunk_size] += probs[j] * window
+                weight_sum[r:r + chunk_size, c:c + chunk_size] += window
+    
+    # Normalize
+    weight_sum = np.maximum(weight_sum, 1e-8)
+    probs_final = probs_sum / weight_sum
+    
+    # Crop to original size
+    probs_final = probs_final[:, half_tile:half_tile + H, half_tile:half_tile + W]
+    
+    # Merge classes if requested - AFTER normalization
+    if merge_clouds:
+        probs_merged = np.zeros((3, H, W), dtype=np.float32)
+        probs_merged[0] = probs_final[0]  # clear
+        probs_merged[1] = probs_final[1] + probs_final[2]  # cloud = thin + thick
+        probs_merged[2] = probs_final[3]  # shadow
+        pred = np.argmax(probs_merged, axis=0).astype(np.uint8)
+    else:
+        pred = np.argmax(probs_final, axis=0).astype(np.uint8)
+    
+    # Apply physical rules
+    if apply_rules:
+        pred = apply_physical_rules(pred, image, merge_clouds, saturation_threshold)
+    
+    return pred
+
+
+# ============================================================================
+# OPTIONAL: EXAMPLE DATA AND VISUALIZATION
+# ============================================================================
+
+def example_data(model_dir: Path, **kwargs):
+    """
+    Load example data for testing (optional function).
+    
+    Returns:
+        Example MSS image as numpy array (4, H, W)
+    """
+    # This is optional - you can provide a small example .npy file
+    example_path = model_dir / "example_mss.npy"
+    
+    if not example_path.exists():
+        # Return synthetic data if no example file
+        print("⚠️  No example data found, generating synthetic")
+        return np.random.rand(4, 512, 512).astype(np.float32) * 0.5
+    
+    return np.load(example_path)
+
+
+def display_results(
+    model_dir: Path,
+    image: np.ndarray,
+    prediction: np.ndarray,
+    stac_item=None,
+    **kwargs
+):
+    """
+    Display prediction results (optional visualization function).
+    
+    Args:
+        model_dir: Model directory
+        image: Input image (4, H, W)
+        prediction: Predicted classes (H, W)
+        stac_item: STAC metadata
+    """
+    try:
+        import matplotlib.pyplot as plt
+        from matplotlib.colors import ListedColormap
+    except ImportError:
+        print("⚠️  matplotlib not installed, skipping visualization")
+        return
+    
+    merge_clouds = prediction.max() <= 2
+    
+    # Color maps
+    if merge_clouds:
+        colors = ['#2E7D32', '#FFFFFF', '#424242']  # clear, cloud, shadow
+        labels = ['Clear', 'Cloud', 'Shadow']
+    else:
+        colors = ['#2E7D32', '#B3E5FC', '#FFFFFF', '#424242']
+        labels = ['Clear', 'Thin Cloud', 'Thick Cloud', 'Shadow']
+    
+    cmap = ListedColormap(colors)
+    
+    # Plot
+    fig, axes = plt.subplots(1, 2, figsize=(12, 5))
+    
+    # RGB composite (use bands 1, 0, 2 as RGB approximation)
+    rgb = np.stack([image[1], image[0], image[2]], axis=-1)
+    rgb = np.clip(rgb * 3, 0, 1)  # Brighten for visibility
+    axes[0].imshow(rgb)
+    axes[0].set_title("MSS RGB Composite")
+    axes[0].axis('off')
+    
+    # Prediction
+    im = axes[1].imshow(prediction, cmap=cmap, vmin=0, vmax=len(labels)-1)
+    axes[1].set_title("Cloud Detection")
+    axes[1].axis('off')
+    
+    # Colorbar
+    cbar = plt.colorbar(im, ax=axes[1], ticks=range(len(labels)))
+    cbar.ax.set_yticklabels(labels)
+    
+    plt.tight_layout()
+    plt.show()

unet.ckpt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:920ac77982e059ba6300f757d5588284cb983f4a5430d05b8103f95101e3470a
+size 154913825

unet.json

@@ -0,0 +1,201 @@
+{
+  "type": "Feature",
+  "stac_version": "1.1.0",
+  "stac_extensions": [
+    "https://stac-extensions.github.io/mlm/v1.5.0/schema.json",
+    "https://stac-extensions.github.io/file/v2.1.0/schema.json"
+  ],
+  "id": "MSS_CLOUDMASK_UNET_EFFB3",
+  "geometry": {
+    "type": "Polygon",
+    "coordinates": [
+      [
+        [
+          -180,
+          -90
+        ],
+        [
+          -180,
+          90
+        ],
+        [
+          180,
+          90
+        ],
+        [
+          180,
+          -90
+        ],
+        [
+          -180,
+          -90
+        ]
+      ]
+    ]
+  },
+  "bbox": [
+    -180,
+    -90,
+    180,
+    90
+  ],
+  "properties": {
+    "datetime": "2026-01-18T22:42:31.441233Z",
+    "created": "2026-01-18T22:42:31.441233Z",
+    "updated": "2026-01-19T01:01:38.488397Z",
+    "title": "MSS Cloud Detection Model (UNet-EfficientNetB3)",
+    "description": "UNet architecture with EfficientNet-B3 encoder for cloud detection in Landsat MSS (Multispectral Scanner) imagery. Trained on CloudSEN12 data emulated to MSS spectral bands using satharmony package. Detects 4 classes: clear, thin cloud, thick cloud, and shadow.",
+    "mlm:name": "mss_cloudmask_unet_effb3",
+    "mlm:architecture": "UNet with EfficientNet-B3 encoder + SCSE attention",
+    "mlm:tasks": [
+      "semantic-segmentation",
+      "cloud-detection"
+    ],
+    "mlm:framework": "pytorch",
+    "mlm:framework_version": "2.5.1+cu121",
+    "mlm:accelerator": "cuda",
+    "mlm:memory_size": 309827650,
+    "mlm:batch_size_suggestion": 8,
+    "mlm:total_parameters": 13223490,
+    "mlm:input": [
+      {
+        "name": "mss_reflectance",
+        "bands": [
+          "Green (500-600nm)",
+          "Red (600-700nm)",
+          "NIR1 (700-800nm)",
+          "NIR2 (800-1100nm)"
+        ],
+        "input": {
+          "shape": [
+            -1,
+            4,
+            "H",
+            "W"
+          ],
+          "dim_order": [
+            "batch",
+            "channel",
+            "height",
+            "width"
+          ],
+          "data_type": "float32"
+        },
+        "norm": {
+          "type": "reflectance",
+          "range": [
+            0.0,
+            1.0
+          ],
+          "description": "TOA reflectance normalized to [0, 1]. DN values should be divided by 10000."
+        },
+        "preprocessing": "Divide DN by 10000 to get reflectance in [0, 1]"
+      }
+    ],
+    "mlm:output": [
+      {
+        "name": "cloud_mask",
+        "classes": [
+          {
+            "id": 0,
+            "name": "clear",
+            "description": "Clear sky"
+          },
+          {
+            "id": 1,
+            "name": "thin_cloud",
+            "description": "Thin/cirrus clouds"
+          },
+          {
+            "id": 2,
+            "name": "thick_cloud",
+            "description": "Thick/opaque clouds"
+          },
+          {
+            "id": 3,
+            "name": "shadow",
+            "description": "Cloud shadow"
+          }
+        ],
+        "result": {
+          "shape": [
+            -1,
+            4,
+            "H",
+            "W"
+          ],
+          "dim_order": [
+            "batch",
+            "class",
+            "height",
+            "width"
+          ],
+          "data_type": "float32"
+        },
+        "description": "Per-pixel logits for 4 classes. Use argmax to get class labels, or softmax for probabilities.",
+        "postprocessing": "Apply argmax(dim=1) to get class labels (0-3), or softmax(dim=1) for probabilities"
+      }
+    ],
+    "mlm:hyperparameters": {
+      "learning_rate": 0.0003,
+      "weight_decay": 0.0001,
+      "optimizer": "AdamW",
+      "scheduler": "CosineAnnealingWarmRestarts",
+      "batch_size": 256,
+      "training_epochs": 55,
+      "final_val_iou": 0.6164,
+      "loss_function": "CrossEntropyLoss",
+      "encoder_depth": 5,
+      "decoder_attention": "SCSE"
+    },
+    "custom:sensor": "Landsat MSS",
+    "custom:spatial_resolution": "60m",
+    "custom:temporal_coverage": "1972-2013",
+    "custom:training_data": "CloudSEN12 emulated to MSS bands",
+    "custom:emulator": "satharmony",
+    "custom:project": "QA4EO-2",
+    "custom:project_url": "https://github.com/IPL-UV/qa4eo",
+    "file:size": 154913825,
+    "dependencies": [
+      "torch>=2.0.0",
+      "pytorch-lightning>=2.0.0",
+      "segmentation-models-pytorch>=0.3.0",
+      "rasterio>=1.3.0",
+      "numpy>=1.21.0"
+    ]
+  },
+  "assets": {
+    "model": {
+      "href": "https://huggingface.co/isp-uv-es/QA4EO-2/resolve/main/unet.ckpt",
+      "type": "application/octet-stream",
+      "title": "PyTorch Lightning checkpoint",
+      "roles": [
+        "mlm:model",
+        "mlm:weights"
+      ],
+      "file:size": 154913825
+    },
+    "load": {
+      "href": "https://huggingface.co/isp-uv-es/QA4EO-2/resolve/main/load.py",
+      "type": "application/x-python-code",
+      "title": "Model loading and inference functions",
+      "roles": [
+        "mlm:inference-code"
+      ]
+    }
+  },
+  "links": [
+    {
+      "rel": "about",
+      "href": "https://github.com/IPL-UV/qa4eo",
+      "type": "text/html",
+      "title": "Project repository"
+    },
+    {
+      "rel": "license",
+      "href": "https://creativecommons.org/licenses/by/4.0/",
+      "type": "text/html",
+      "title": "CC-BY-4.0"
+    }
+  ]
+}