Automating Cloud Removal in Sentinel-2 Time Series

import os
import glob
import numpy as np
import rasterio
import xarray as xr
from pathlib import Path
from typing import List, Optional, Dict


def _extract_date(filename: str) -> str:
    """Return the 8-digit acquisition date (YYYYMMDD) from a Sentinel-2 filename."""
    stem = Path(filename).stem           # e.g. T32TQM_20230615T105021_SCL
    parts = stem.split("_")
    for part in parts:
        if len(part) >= 8 and part[:8].isdigit():
            return part[:8]
    raise ValueError(f"Cannot parse date from filename: {filename}")


def collect_file_pairs(
    bands_dir: str,
    scl_dir: str,
    target_bands: List[str]
) -> Dict[str, Dict]:
    """
    Return a dict keyed by YYYYMMDD, each value a dict with 'scl' path and
    per-band paths.  Dates missing any band or SCL file are silently skipped.
    """
    scl_files = sorted(glob.glob(os.path.join(scl_dir, "*_SCL*.tif")))
    pairs: Dict[str, Dict] = {}

    for scl_path in scl_files:
        date = _extract_date(scl_path)
        band_paths = {}
        for band in target_bands:
            matches = glob.glob(
                os.path.join(bands_dir, f"*{date}*_{band}*.tif")
            )
            if not matches:
                break                     # skip incomplete dates
            band_paths[band] = matches[0]
        else:
            pairs[date] = {"scl": scl_path, "bands": band_paths}

    if not pairs:
        raise FileNotFoundError(
            "No complete date entries found.  Check directory paths and filename patterns."
        )
    return pairs

VALID_SCL_CLASSES = [4, 5, 6]   # adjust to include class 7 in persistently cloudy areas


def build_scl_mask(scl_path: str, valid_classes: List[int] = VALID_SCL_CLASSES) -> np.ndarray:
    """
    Read the SCL band and return a float32 array:
      1.0  → valid pixel
      NaN  → cloud / shadow / invalid pixel
    """
    with rasterio.open(scl_path) as src:
        scl = src.read(1)                  # SCL is always single-band

    mask = np.where(np.isin(scl, valid_classes), 1.0, np.nan).astype(np.float32)

    # Sanity check: a mask that is entirely NaN means the scene is fully clouded
    valid_fraction = np.nansum(mask) / mask.size
    if valid_fraction < 0.01:
        import warnings
        warnings.warn(
            f"SCL mask for {Path(scl_path).name} is >99 % invalid "
            f"(valid fraction: {valid_fraction:.3f}). "
            "Consider excluding this date."
        )
    return mask

SCALE_FACTOR = 0.0001   # ESA sen2cor Level-2A scaling


def load_masked_bands(
    band_paths: Dict[str, str],
    mask: np.ndarray
) -> Dict[str, np.ndarray]:
    """
    Load each band, scale to surface reflectance, and apply the validity mask.
    Returns a dict of float32 arrays with NaN on invalid pixels.
    """
    arrays: Dict[str, np.ndarray] = {}
    for band_name, path in band_paths.items():
        with rasterio.open(path) as src:
            raw = src.read(1).astype(np.float32)
            # Preserve spatial profile for export (taken from the last band read — reuse later)
            profile = src.profile
        reflectance = np.clip(raw * SCALE_FACTOR, 0.0, 1.0)
        arrays[band_name] = reflectance * mask   # NaN where mask is NaN
    return arrays, profile

def build_composite(
    file_pairs: Dict[str, Dict],
    target_bands: List[str],
    valid_classes: List[int] = VALID_SCL_CLASSES
) -> tuple[xr.Dataset, dict]:
    """
    Iterate over all dates, mask each date's bands, stack into an xarray Dataset,
    and return a median composite along the time axis.
    """
    date_datasets = []
    dates = sorted(file_pairs.keys())       # chronological order
    profile_ref = None

    for date in dates:
        scl_path = file_pairs[date]["scl"]
        band_paths = file_pairs[date]["bands"]

        mask = build_scl_mask(scl_path, valid_classes)
        arrays, profile = load_masked_bands(band_paths, mask)

        if profile_ref is None:
            profile_ref = profile

        ds = xr.Dataset(
            {band: xr.DataArray(arr, dims=["y", "x"])
             for band, arr in arrays.items()}
        )
        date_datasets.append(ds)

    if not date_datasets:
        raise ValueError("All dates were skipped (fully clouded or missing files).")

    # Stack along time dimension
    stack = xr.concat(date_datasets, dim="time")
    stack = stack.assign_coords(time=dates)

    # Observation-count layer: how many valid (non-NaN) pixels contributed per position
    obs_count = (~stack.isnull()).sum(dim="time")

    # Median composite — NaN-aware
    composite = stack.median(dim="time", skipna=True)

    # Flag pixels with fewer than 2 valid observations as unreliable
    min_obs = 2
    for band in target_bands:
        composite[band] = xr.where(obs_count[band] >= min_obs, composite[band], np.nan)

    return composite, profile_ref

def export_composite(
    composite: xr.Dataset,
    profile: dict,
    out_path: str,
    target_bands: List[str]
) -> None:
    """Write the cloud-free composite to a multi-band GeoTIFF."""
    profile = profile.copy()
    profile.update(
        count=len(target_bands),
        dtype="float32",
        nodata=np.nan,
        compress="deflate",
        predictor=3,        # floating-point predictor for better compression
        tiled=True,
        blockxsize=512,
        blockysize=512,
    )

    with rasterio.open(out_path, "w", **profile) as dst:
        for i, band in enumerate(target_bands, start=1):
            arr = composite[band].values
            # Validate before writing
            assert arr.dtype == np.float32, f"Band {band} dtype mismatch"
            valid_vals = arr[~np.isnan(arr)]
            assert valid_vals.min() >= 0.0 and valid_vals.max() <= 1.0, (
                f"Band {band} contains out-of-range reflectance values: "
                f"min={valid_vals.min():.4f}, max={valid_vals.max():.4f}"
            )
            dst.write(arr, i)

    print(f"Composite written to {out_path}")
    print(f"  Bands: {target_bands}")
    print(f"  Valid pixel coverage: "
          f"{np.isfinite(composite[target_bands[0]].values).mean() * 100:.1f} %")


# --- Entry point ---
if __name__ == "__main__":
    TARGET_BANDS = ["B04", "B08", "B11"]   # Red, NIR, SWIR — adjust for your index needs

    pairs = collect_file_pairs(
        bands_dir="./S2_L2A/bands",
        scl_dir="./S2_L2A/SCL",
        target_bands=TARGET_BANDS,
    )
    composite, profile = build_composite(pairs, TARGET_BANDS)
    export_composite(composite, profile, "./output/cloud_free_composite.tif", TARGET_BANDS)

import rasterio
import numpy as np

with rasterio.open("./output/cloud_free_composite.tif") as src:
    arr = src.read(1)                  # first band (B04)
    assert src.crs is not None, "Output CRS missing"
    valid = arr[~np.isnan(arr)]
    print(f"B04 median reflectance: {np.median(valid):.4f}")   # expect ~0.05–0.15 for vegetation
    assert valid.min() >= 0.0 and valid.max() <= 1.0
    print("Validation passed.")