Vector vs Raster Annotation Workflows for Geospatial ML

import rasterio
import json
from pathlib import Path

def extract_tile_metadata(tif_path: str) -> dict:
    """Return the spatial metadata needed to pin annotation CRS and rasterize correctly."""
    with rasterio.open(tif_path) as src:
        return {
            "path": tif_path,
            "crs_wkt": src.crs.to_wkt(),
            "epsg": src.crs.to_epsg(),
            "transform": list(src.transform),   # 6-element affine tuple
            "width": src.width,
            "height": src.height,
            "dtype": src.dtypes[0],
            "count": src.count,
        }

# Usage
meta = extract_tile_metadata("/data/tiles/tile_0042.tif")
Path("/data/tiles/tile_0042_meta.json").write_text(json.dumps(meta, indent=2))

import geopandas as gpd
from shapely.validation import make_valid
import logging

logging.basicConfig(level=logging.INFO, format="%(levelname)s %(message)s")

def validate_vector_annotations(
    input_path: str,
    output_path: str,
    target_epsg: int = 32618,
) -> gpd.GeoDataFrame:
    """
    Load, reproject, repair, and export a vector annotation file.

    Args:
        input_path:   Path to the raw GeoJSON or GeoPackage from the annotation tool.
        output_path:  Destination GeoJSON path for validated output.
        target_epsg:  EPSG code matching the source imagery CRS (default UTM 18N).
    Returns:
        Validated GeoDataFrame ready for rasterization or ML export.
    """
    gdf = gpd.read_file(input_path)
    logging.info("Loaded %d features. Source CRS: %s", len(gdf), gdf.crs)

    # --- CRS enforcement ---
    target_crs = f"EPSG:{target_epsg}"
    if gdf.crs is None:
        raise ValueError("Input file has no CRS. Declare it explicitly before export.")
    if gdf.crs.to_epsg() != target_epsg:
        logging.warning("Reprojecting from %s to %s", gdf.crs, target_crs)
        gdf = gdf.to_crs(target_crs)

    # --- Topology repair ---
    invalid_mask = ~gdf.geometry.is_valid
    n_invalid = int(invalid_mask.sum())
    if n_invalid:
        logging.warning("Repairing %d invalid geometries via make_valid().", n_invalid)
        gdf.loc[invalid_mask, "geometry"] = (
            gdf.loc[invalid_mask, "geometry"].apply(make_valid)
        )

    # --- Drop nulls and empty rings post-repair ---
    before = len(gdf)
    gdf = gdf[gdf.geometry.notna() & ~gdf.geometry.is_empty].copy()
    logging.info("Dropped %d null/empty geometries.", before - len(gdf))

    # --- Required attribute schema ---
    for col in ("class_id", "class_name", "annotator_id"):
        if col not in gdf.columns:
            raise ValueError(f"Missing required column: {col}")

    gdf.to_file(output_path, driver="GeoJSON")
    logging.info("Exported %d valid features → %s", len(gdf), output_path)
    return gdf

import rasterio
from rasterio.features import rasterize
from rasterio.transform import Affine
import numpy as np
import geopandas as gpd
import logging

def generate_raster_mask(
    vector_path: str,
    template_tif: str,
    output_mask: str,
    class_mapping: dict[str, int],
) -> np.ndarray:
    """
    Rasterize validated vector annotations against a source imagery tile.

    Args:
        vector_path:   Path to validated GeoJSON.
        template_tif:  Source GeoTIFF whose grid pins the output mask.
        output_mask:   Output GeoTIFF mask path.
        class_mapping: Dict mapping class_name strings to uint8 IDs.
    Returns:
        The rasterized mask array (H, W) for downstream assertion.
    """
    with rasterio.open(template_tif) as src:
        transform: Affine = src.transform
        out_shape: tuple[int, int] = (src.height, src.width)
        crs = src.crs
        logging.info("Template: shape=%s  CRS=%s", out_shape, crs)

    gdf = gpd.read_file(vector_path)
    if gdf.crs.to_epsg() != crs.to_epsg():
        logging.warning("Reprojecting annotations to match tile CRS.")
        gdf = gdf.to_crs(crs)

    gdf["class_id"] = (
        gdf["class_name"].map(class_mapping).fillna(0).astype(np.uint8)
    )

    # Sort descending by class_id so higher-priority classes overwrite lower ones
    gdf = gdf.sort_values("class_id", ascending=False)
    shapes = (
        (geom, int(val))
        for geom, val in zip(gdf.geometry, gdf["class_id"])
        if geom is not None and not geom.is_empty
    )

    mask = rasterize(
        shapes,
        out_shape=out_shape,
        transform=transform,
        fill=0,
        dtype=np.uint8,
        merge_alg=rasterio.enums.MergeAlg.replace,
    )

    with rasterio.open(
        output_mask, "w",
        driver="GTiff",
        height=out_shape[0],
        width=out_shape[1],
        count=1,
        dtype=np.uint8,
        crs=crs,
        transform=transform,
        compress="lzw",
    ) as dst:
        dst.write(mask, 1)

    logging.info("Mask → %s  unique classes: %s", output_mask, np.unique(mask).tolist())
    return mask

import torch
from torch.utils.data import Dataset
import rasterio
from rasterio.features import rasterize
import geopandas as gpd
import numpy as np
from pathlib import Path

class GeoAnnotationDataset(Dataset):
    """Lazy-rasterize vector annotations to uint8 masks per tile."""

    def __init__(
        self,
        tile_dir: str,
        annotation_dir: str,
        class_mapping: dict[str, int],
        transform=None,
    ):
        self.tiles = sorted(Path(tile_dir).glob("*.tif"))
        self.annotation_dir = Path(annotation_dir)
        self.class_mapping = class_mapping
        self.transform = transform

    def __len__(self) -> int:
        return len(self.tiles)

    def __getitem__(self, idx: int) -> dict:
        tif_path = self.tiles[idx]
        geojson_path = self.annotation_dir / tif_path.with_suffix(".geojson").name

        with rasterio.open(tif_path) as src:
            image = src.read().astype(np.float32)   # (C, H, W)
            t, shape, crs = src.transform, src.shape, src.crs

        gdf = gpd.read_file(geojson_path)
        if gdf.crs.to_epsg() != crs.to_epsg():
            gdf = gdf.to_crs(crs)

        gdf["class_id"] = gdf["class_name"].map(self.class_mapping).fillna(0).astype(np.uint8)
        shapes = ((g, int(v)) for g, v in zip(gdf.geometry, gdf["class_id"]) if g and not g.is_empty)
        mask = rasterize(shapes, out_shape=shape, transform=t, fill=0, dtype=np.uint8)

        sample = {"image": torch.from_numpy(image), "mask": torch.from_numpy(mask.astype(np.int64))}
        if self.transform:
            sample = self.transform(sample)
        return sample

# label_studio_export_hook.py
import subprocess
import sys

def on_export_complete(export_path: str, target_epsg: int = 32618) -> None:
    """Called by Label Studio export webhook — validates before committing to dataset."""
    result = subprocess.run(
        ["python", "validate_annotations.py", export_path, str(target_epsg)],
        capture_output=True, text=True,
    )
    if result.returncode != 0:
        raise RuntimeError(f"Annotation validation failed:\n{result.stderr}")
    print(f"Validation passed: {export_path}")

# dvc.yaml
stages:
  validate_vectors:
    cmd: python validate_annotations.py data/raw_annotations data/validated_annotations
    deps:
      - data/raw_annotations
      - validate_annotations.py
    outs:
      - data/validated_annotations

  rasterize_masks:
    cmd: python rasterize_pipeline.py data/validated_annotations data/tiles data/masks
    deps:
      - data/validated_annotations
      - data/tiles
      - rasterize_pipeline.py
    outs:
      - data/masks

import geopandas as gpd
import rasterio
import numpy as np
from pathlib import Path

def assert_annotation_integrity(
    geojson_path: str,
    mask_path: str,
    template_tif: str,
    class_mapping: dict[str, int],
) -> None:
    """Raise AssertionError on any integrity violation."""

    # 1. All vector geometries valid
    gdf = gpd.read_file(geojson_path)
    assert gdf.geometry.is_valid.all(), "Invalid geometries detected — run make_valid()"

    # 2. CRS roundtrip: vector CRS matches tile CRS
    with rasterio.open(template_tif) as src:
        tile_epsg = src.crs.to_epsg()
        tile_transform = src.transform
        tile_shape = src.shape
    assert gdf.crs.to_epsg() == tile_epsg, (
        f"CRS mismatch: annotations={gdf.crs.to_epsg()} tile={tile_epsg}"
    )

    # 3. Mask shape matches tile shape
    with rasterio.open(mask_path) as msrc:
        mask = msrc.read(1)
        assert msrc.crs.to_epsg() == tile_epsg, "Mask CRS mismatch"
        assert msrc.transform == tile_transform, "Mask affine transform mismatch"
    assert mask.shape == tile_shape, f"Mask shape {mask.shape} ≠ tile shape {tile_shape}"

    # 4. Mask contains only known class IDs
    valid_ids = {0, 255} | set(class_mapping.values())
    observed_ids = set(np.unique(mask).tolist())
    unknown = observed_ids - valid_ids
    assert not unknown, f"Mask contains unknown class IDs: {unknown}"

    # 5. Non-empty annotation coverage
    coverage = float((mask > 0).sum()) / mask.size
    assert coverage > 0.001, f"Mask is nearly empty (coverage={coverage:.4f})"

    print(f"All integrity checks passed. Coverage: {coverage:.2%}")

# .pre-commit-config.yaml
- repo: local
  hooks:
    - id: annotation-integrity
      name: Annotation integrity check
      entry: python scripts/assert_annotation_integrity.py
      language: python
      files: \.geojson$