deepcontrail-segmentation 0.1.0__py3-none-any.whl

deepcontrail/segmentation/__init__.py

@@ -0,0 +1,51 @@
+"""Semantic segmentation of contrails in satellite imagery.
+
+The data pipeline imported here is torch-free; ``models``, ``losses`` and
+``training`` are subpackages that need the training stack. ``train`` (from
+``training``), ``load_pretrained`` and ``predict`` (from ``pretrained``) are loaded
+on first access, so this import stays light.
+"""
+
+from importlib.metadata import version
+from typing import TYPE_CHECKING, Any
+
+from deepcontrail.segmentation.config import (
+    AugmentationConfig,
+    NormalizeConfig,
+    SegmentationDataConfig,
+)
+from deepcontrail.segmentation.dataset import SegmentationDataset
+
+__version__ = version("deepcontrail-segmentation")
+
+if TYPE_CHECKING:
+    from deepcontrail.segmentation.pretrained import load_pretrained, predict
+    from deepcontrail.segmentation.training.run import train
+
+_LAZY = {
+    "train": "deepcontrail.segmentation.training.run",
+    "load_pretrained": "deepcontrail.segmentation.pretrained",
+    "predict": "deepcontrail.segmentation.pretrained",
+}
+
+
+def __getattr__(name: str) -> Any:  # noqa: ANN401
+    """Load ``train``, ``load_pretrained`` and ``predict`` lazily: they import torch."""
+    if name in _LAZY:
+        import importlib
+
+        return getattr(importlib.import_module(_LAZY[name]), name)
+    msg = f"module {__name__!r} has no attribute {name!r}"
+    raise AttributeError(msg)
+
+
+__all__ = [
+    "AugmentationConfig",
+    "NormalizeConfig",
+    "SegmentationDataConfig",
+    "SegmentationDataset",
+    "__version__",
+    "load_pretrained",
+    "predict",
+    "train",
+]

deepcontrail/segmentation/config.py

@@ -0,0 +1,94 @@
+"""Pydantic configuration of the segmentation data pipeline.
+
+Uses only pydantic so that ``import deepcontrail.segmentation`` stays torch-free.
+"""
+
+from __future__ import annotations
+
+from typing import Annotated, Literal
+
+from pydantic import BaseModel, ConfigDict, Field
+
+from deepcontrail.core.config import DatasetConfig
+
+
+class NormalizeConfig(BaseModel):
+    """Channel-wise normalization statistics.
+
+    Defaults to ImageNet mean/std, the standard choice when fine-tuning
+    pretrained encoders (ResNet, EfficientNet, etc.).
+
+    For model-specific normalization, training pipelines should query
+    ``smp.encoders.get_preprocessing_params(encoder_name)`` and pass the
+    resulting ``mean``/``std`` here to match the encoder's expectations.
+    """
+
+    model_config = ConfigDict(frozen=True, extra="forbid")
+
+    mean: tuple[float, float, float] = (0.485, 0.456, 0.406)
+    std: tuple[float, float, float] = (0.229, 0.224, 0.225)
+
+
+class AugmentationConfig(BaseModel):
+    """Spatial and photometric augmentation probabilities.
+
+    Each probability defaults to ``0.0`` (disabled) except the three spatial
+    flips/rotates, so existing YAML configs without the new fields keep their
+    previous behaviour. Photometric transforms (``gaussian_blur``,
+    ``brightness_contrast``) apply to the image only — albumentations
+    co-transforms the mask for the spatial ones.
+    """
+
+    model_config = ConfigDict(frozen=True, extra="forbid")
+
+    # Spatial — applied to both image and mask.
+    horizontal_flip: Annotated[float, Field(ge=0.0, le=1.0)] = 0.5
+    vertical_flip: Annotated[float, Field(ge=0.0, le=1.0)] = 0.5
+    random_rotate_90: Annotated[float, Field(ge=0.0, le=1.0)] = 0.5
+    shift_scale_rotate: Annotated[float, Field(ge=0.0, le=1.0)] = 0.0
+    shift_limit: float = 0.0625
+    scale_limit: float = 0.1
+    rotate_limit: int = 45
+    # Photometric / pixel-level — image only.
+    gaussian_blur: Annotated[float, Field(ge=0.0, le=1.0)] = 0.0
+    blur_limit: tuple[int, int] = (3, 7)
+    brightness_contrast: Annotated[float, Field(ge=0.0, le=1.0)] = 0.0
+
+
+class SegmentationDataConfig(DatasetConfig):
+    """Complete segmentation data pipeline configuration.
+
+    Extends ``DatasetConfig`` (root, split, image construction, pixel shift and
+    train/val partition) with the label, resize, augmentation, normalization
+    and DataLoader parameters (batch_size, num_workers), so one flat ``data:``
+    YAML section reproduces a dataset setup.
+
+    Parameters
+    ----------
+    mask_type
+        Which label artifact to load. ``"binary"`` reads ``semantic_masks/``
+        PNGs (strict consensus ground truth). ``"probability"`` reads
+        ``probability_masks/`` and returns float32 in ``[0, 1]`` for
+        soft-label training.
+    image_size
+        Expected spatial size of images (height and width).
+    batch_size
+        Batch size for the DataLoader (not used by the Dataset itself).
+    num_workers
+        Number of DataLoader workers (not used by the Dataset itself).
+    augmentation
+        Augmentation config for training, or ``None`` for no augmentation.
+    normalize
+        Channel-wise normalization statistics.
+    """
+
+    # Unlike the shared core DatasetConfig, the segmentation data section is complete here, so a
+    # misspelt key is an error instead of a silently ignored setting.
+    model_config = ConfigDict(frozen=True, extra="forbid")
+
+    mask_type: Literal["binary", "probability"] = "binary"
+    image_size: Annotated[int, Field(gt=0)] = 256
+    batch_size: Annotated[int, Field(gt=0)] = 16
+    num_workers: Annotated[int, Field(ge=0)] = 4
+    augmentation: AugmentationConfig | None = None
+    normalize: NormalizeConfig = NormalizeConfig()

deepcontrail/segmentation/dataset.py

@@ -0,0 +1,275 @@
+"""Semantic segmentation dataset over a DeepContrail-format dataset root.
+
+Builds each image from ``spectral_bands/`` with the configured false-colour recipe
+(band keys from the root's ``bands.json``) and pairs it with its mask.
+``import deepcontrail.segmentation`` requires only numpy + Pillow — torch is
+imported lazily inside ``__getitem__`` and is optional.
+"""
+
+from __future__ import annotations
+
+import warnings
+from pathlib import Path
+from typing import TYPE_CHECKING, Any
+
+import numpy as np
+
+from deepcontrail.core.contract import read_band_manifest, read_dataset_info
+from deepcontrail.core.false_color import check_false_color, load_spectral_false_color
+from deepcontrail.core.readers import (
+    binarize_mask,
+    load_mask,
+    shift_image,
+    to_numpy,
+    validate_binary_mask,
+)
+from deepcontrail.core.splits import context_frames, partition_stems
+from deepcontrail.segmentation.tensors import to_tensors, to_tensors_soft
+
+if TYPE_CHECKING:
+    import albumentations as A
+
+    from deepcontrail.segmentation.config import SegmentationDataConfig
+
+
+class SegmentationDataset:
+    """Per-sample semantic segmentation dataset over a DeepContrail-format root.
+
+    Parameters
+    ----------
+    config
+        A ``SegmentationDataConfig`` specifying root path, split, augmentation, and
+        normalization parameters.
+
+    Raises
+    ------
+    DatasetContractError
+        If the root's ``bands.json`` or ``dataset_info.json`` is missing or invalid,
+        or if ``false_color`` uses a band key the manifest does not declare.
+    ValueError
+        If any labelled image in the split directory has no matching mask.
+    ImportError
+        If augmentation or resize is requested but albumentations is not
+        installed.
+    """
+
+    def __init__(self, config: SegmentationDataConfig) -> None:
+        self.config = config
+        self._root = Path(config.root)
+        self.manifest = read_band_manifest(self._root)
+        check_false_color(config.false_color, self.manifest)
+        self._warn_on_pixel_shift_mismatch()
+
+        img_subdir = "spectral_bands"
+        mask_subdir = "probability_masks" if config.mask_type == "probability" else "semantic_masks"
+        ds = "test" if config.split == "test" else "train"
+        self._dir_pairs = [
+            (self._root / img_subdir / ds, self._root / mask_subdir / ds),
+        ]
+        self._image_dir, self._mask_dir = self._dir_pairs[0]
+
+        self._stem_dirs: dict[str, tuple[Path, Path]] = {}
+        self.stems = self._discover_stems()
+        self.native_size = self._native_size()
+        self._aug_transform, self._resize_transform = self._build_transform()
+        # DataLoader worker the augmentation was last reseeded for (None: main process).
+        self._augmentation_worker: int | None = None
+
+    def _reseed_augmentation_in_worker(self) -> None:
+        """Give each DataLoader worker its own augmentation stream.
+
+        Every worker starts with a copy of the dataset, generators included, so without this
+        all workers would draw the same augmentations. The first item a worker loads reseeds the
+        copy from PyTorch's per-worker seed, which is reproducible under a fixed global seed and
+        changes each epoch unless the workers persist. The main process keeps the
+        configuration's seed.
+        """
+        try:
+            from torch.utils.data import get_worker_info
+        except ImportError:  # without torch there are no DataLoader workers
+            return
+        info = get_worker_info()
+        if info is None or info.id == self._augmentation_worker:
+            return
+        from deepcontrail.segmentation.transforms import seed_augmentation
+
+        seed_augmentation(self._aug_transform, info.seed)  # type: ignore[arg-type]
+        self._augmentation_worker = info.id
+
+    def _warn_on_pixel_shift_mismatch(self) -> None:
+        """Warn when ``pixel_shift`` differs from the dataset's recommendation."""
+        recommended = read_dataset_info(self._root).recommended_pixel_shift
+        if recommended is not None and recommended != self.config.pixel_shift:
+            msg = (
+                f"pixel_shift={self.config.pixel_shift} but {self._root} recommends "
+                f"{recommended} (metadata/dataset_info.json): images and masks may be misaligned"
+            )
+            warnings.warn(msg, UserWarning, stacklevel=3)
+
+    @property
+    def native_resolution(self) -> int:
+        """Native square resolution of the images (256 for GOES-16)."""
+        return self.native_size
+
+    def _discover_stems(self) -> list[str]:
+        """Scan the image directory and partition stems for the requested split.
+
+        ``"test"`` returns every stem in ``test/``.  ``"train"`` and ``"val"``
+        both read from ``train/`` and split deterministically:
+
+        * When ``fold`` is set, k-fold CV partitioning (modular arithmetic)
+          takes precedence.
+        * Otherwise ``val_fraction`` determines the split — stems are
+          shuffled with ``seed`` and the first ``val_fraction * N`` go to
+          validation, the rest to training.
+        """
+        all_stems: list[str] = []
+        for img_dir, mask_dir in self._dir_pairs:
+            scanned = self._scan_stems(img_dir, mask_dir)
+            for stem in scanned:
+                self._stem_dirs[stem] = (img_dir, mask_dir)
+            all_stems.extend(scanned)
+        return partition_stems(
+            all_stems,
+            self.config.split,
+            seed=self.config.seed,
+            val_fraction=self.config.val_fraction,
+            fold=self.config.fold,
+            num_folds=self.config.num_folds,
+        )
+
+    def _scan_stems(self, image_dir: Path, mask_dir: Path) -> list[str]:
+        """Scan ``image_dir`` for paired images/masks, return sorted stems.
+
+        Checks that each image has a matching mask file; only the context frames of a dataset
+        that declares them (``temporal: true``) may have none. For deeper validation
+        (corruption, sizes, mask values), run ``deepcontrail-validate <root> --deep`` on the
+        dataset root.
+        """
+        mask_ext = ".npy" if self.config.mask_type == "probability" else ".png"
+        image_stems = [p.stem for p in sorted(image_dir.glob("*.npy"))]
+        context = context_frames(image_stems, temporal=read_dataset_info(self._root).temporal)
+        stems: list[str] = []
+        missing: list[str] = []
+
+        for stem in image_stems:
+            mask_path = mask_dir / f"{stem}{mask_ext}"
+            if mask_path.exists():
+                stems.append(stem)
+            elif stem not in context:
+                missing.append(f"{stem}: no mask (expected {mask_path})")
+
+        if missing:
+            msg = f"Dataset validation found {len(missing)} missing mask(s):\n"
+            msg += "\n".join(f"  - {e}" for e in missing)
+            raise ValueError(msg)
+
+        return stems
+
+    def _native_size(self) -> int:
+        """Native square resolution of the images (256 for GOES-16)."""
+        if not self.stems:
+            return self.config.image_size
+        arr = np.load(self._image_dir / f"{self.stems[0]}.npy", mmap_mode="r")
+        return arr.shape[0]
+
+    def _needs_resize(self) -> bool:
+        """Whether images must be resized to match ``image_size``.
+
+        Compares the native resolution against ``image_size``; the mask is
+        always returned at native resolution regardless.
+        """
+        if not self.stems:
+            return False
+        return self.native_size != self.config.image_size
+
+    def _build_transform(
+        self,
+    ) -> tuple[A.Compose | None, A.Resize | None]:
+        """Build the augmentation and (optional) image-only resize transforms.
+
+        Returns
+        -------
+        tuple[A.Compose | None, A.Resize | None]
+            ``(augmentation, image_resize)``. ``augmentation`` co-transforms
+            the image and mask at native resolution; ``image_resize`` resizes
+            only the image, when ``image_size`` differs from native.
+        """
+        has_augmentation = self.config.augmentation is not None
+        needs_resize = self._needs_resize()
+
+        if not has_augmentation and not needs_resize:
+            return None, None
+
+        try:
+            from deepcontrail.segmentation.transforms import build_augmentation, build_image_resize
+        except ImportError:
+            msg = (
+                "albumentations is required for augmentation and resizing. "
+                "Install it with: pip install deepcontrail[train]"
+            )
+            raise ImportError(msg) from None
+
+        augmentation = (
+            build_augmentation(self.config.augmentation, seed=self.config.seed)  # type: ignore[arg-type]
+            if has_augmentation
+            else None
+        )
+        image_resize = build_image_resize(self.config.image_size) if needs_resize else None
+        return augmentation, image_resize
+
+    def __len__(self) -> int:
+        """Return the number of image/mask pairs in this split."""
+        return len(self.stems)
+
+    def __getitem__(self, index: int) -> tuple[Any, Any]:
+        """Load and return ``(image, mask)`` for the given index.
+
+        With torch installed, returns:
+        - ``image``: float32 tensor, shape ``(C, H, W)``, ImageNet-normalized.
+        - ``mask``: long tensor, shape ``(H, W)``, values in ``{0, 1}``.
+
+        Without torch, returns numpy arrays:
+        - ``image``: float32 ndarray, shape ``(H, W, C)``, in ``[0, 1]``.
+        - ``mask``: uint8 ndarray, shape ``(H, W)``, values in ``{0, 1}``.
+        """
+        stem = self.stems[index]
+
+        img_dir, mask_dir = self._stem_dirs.get(stem, (self._image_dir, self._mask_dir))
+        image = load_spectral_false_color(
+            img_dir / f"{stem}.npy", self.config.false_color, self.manifest
+        )
+        if self.config.pixel_shift != 0.0:
+            image = shift_image(image, self.config.pixel_shift)
+        if self.config.mask_type == "probability":
+            mask = np.load(mask_dir / f"{stem}.npy")
+        else:
+            mask = load_mask(mask_dir / f"{stem}.png")
+            mask = binarize_mask(mask)
+
+        assert image.shape[:2] == mask.shape, (
+            f"spatial mismatch at native resolution: image {image.shape[:2]} vs mask {mask.shape}"
+        )
+
+        # Augmentation runs at native resolution on BOTH image and mask so they
+        # stay aligned and the mask is never upscaled by interpolation.
+        if self._aug_transform is not None:
+            self._reseed_augmentation_in_worker()
+            transformed = self._aug_transform(image=image, mask=mask)
+            image = transformed["image"]
+            mask = transformed["mask"]
+
+        # The image alone is resized to image_size; the mask stays native.
+        if self._resize_transform is not None:
+            image = self._resize_transform(image=image)["image"]
+
+        if self.config.mask_type != "probability":
+            validate_binary_mask(mask)
+
+        try:
+            import torch as _  # noqa: F401
+        except ImportError:
+            return to_numpy(image, mask)
+        if self.config.mask_type == "probability":
+            return to_tensors_soft(image, mask, self.config.normalize)
+        return to_tensors(image, mask, self.config.normalize)

deepcontrail/segmentation/losses/__init__.py

@@ -0,0 +1,21 @@
+"""Differentiable training losses for segmentation."""
+
+from deepcontrail.segmentation.losses.config import LossConfig
+from deepcontrail.segmentation.losses.losses import (
+    BCELoss,
+    CompositeLoss,
+    DiceLoss,
+    FocalLoss,
+    TverskyLoss,
+    build_loss,
+)
+
+__all__ = [
+    "BCELoss",
+    "CompositeLoss",
+    "DiceLoss",
+    "FocalLoss",
+    "LossConfig",
+    "TverskyLoss",
+    "build_loss",
+]

deepcontrail/segmentation/losses/config.py

@@ -0,0 +1,87 @@
+"""Pydantic configuration model for loss functions.
+
+Serializable via ``model_dump()`` for MLflow logging.
+Uses only pydantic — no torch dependency.
+"""
+
+from __future__ import annotations
+
+from typing import Annotated, Any, Literal
+
+from pydantic import BaseModel, ConfigDict, Field, model_validator
+
+COMPOSITE_COMPONENTS = ("dice", "tversky", "bce", "focal")
+
+
+class LossConfig(BaseModel):
+    """Loss function configuration for training.
+
+    Parameters
+    ----------
+    name
+        Loss function name: ``"dice"``, ``"tversky"``, ``"bce"``,
+        ``"focal"``, or ``"composite"``.
+    smooth
+        Smoothing factor of the Dice loss. The Tversky loss always uses 0.
+    alpha
+        Tversky alpha (weight for false positives).
+    beta
+        Tversky beta (weight for false negatives).
+    gamma
+        Tversky focusing parameter.
+    from_logits
+        Whether the Dice and Tversky losses receive raw logits (sigmoid applied internally). The
+        BCE and focal losses always take logits.
+    composite_weights
+        Weights of the composite loss components, among ``"dice"``, ``"tversky"``, ``"bce"`` and
+        ``"focal"``: none negative, at least one positive.
+        Keys may be ``"dice"``, ``"tversky"``, ``"bce"``, ``"focal"``.
+        Only used when ``name="composite"``.
+    pos_weight
+        BCE positive-class weight.  Set to ``num_bg / num_fg`` to compensate
+        for class imbalance (e.g. ``175.0`` for 0.567 % foreground).
+        ``None`` means no weighting.
+    smooth_factor
+        BCE label-smoothing factor.  ``None`` means no smoothing.
+    focal_gamma
+        Focal loss focusing parameter — down-weights easy (confident) examples.
+        ``gamma=0`` is equivalent to BCE; ``gamma=2`` is the standard value.
+    focal_alpha
+        Focal loss class-balance factor.  ``None`` means no balancing.
+    """
+
+    model_config = ConfigDict(frozen=True, extra="forbid")
+
+    name: Literal["dice", "tversky", "bce", "focal", "composite"] = "dice"
+    smooth: Annotated[float, Field(ge=0.0)] = 1.0
+    alpha: Annotated[float, Field(ge=0.0, le=1.0)] = 0.7
+    beta: Annotated[float, Field(ge=0.0, le=1.0)] = 0.3
+    gamma: Annotated[float, Field(gt=0.0)] = 1.0
+    from_logits: bool = True
+    composite_weights: dict[str, float] = {"dice": 0.5, "tversky": 0.5}
+    pos_weight: float | None = None
+    smooth_factor: float | None = None
+    focal_gamma: Annotated[float, Field(ge=0.0)] = 2.0
+    focal_alpha: float | None = None
+
+    @model_validator(mode="after")
+    def _check_composite_weights(self) -> LossConfig:
+        """Composite components must be known, non-negative, and at least one positive."""
+        unknown = sorted(set(self.composite_weights) - set(COMPOSITE_COMPONENTS))
+        if unknown:
+            expected = list(COMPOSITE_COMPONENTS)
+            msg = f"unknown composite_weights {unknown}; expected some of {expected}"
+            raise ValueError(msg)
+        negative = sorted(k for k, w in self.composite_weights.items() if w < 0)
+        if negative:
+            raise ValueError(f"composite_weights must not be negative: {negative}")
+        if self.name == "composite" and not any(w > 0 for w in self.composite_weights.values()):
+            raise ValueError(
+                "a composite loss needs at least one positive weight in composite_weights"
+            )
+        return self
+
+    @classmethod
+    def from_dict(cls, d: dict[str, Any]) -> LossConfig:
+        """Build a ``LossConfig`` from a plain dictionary."""
+        return cls.model_validate(d)