Add valid LiDAR mask support to tiling workflow

Author: LucaRom

geo_deep_learning/tools/water_extraction/README.md

@@ -57,6 +57,7 @@ AOI_Name/
 ├── dtm.tif           # Digital Terrain Model (bare earth elevation)
 ├── dsm.tif           # Digital Surface Model (first-return elevation)
 ├── intensity.tif     # (Optional) LiDAR intensity
+├── valid_lidar_mask.gpkg  # (Optional) polygon coverage of valid LiDAR regions
 ├── aoi.shp           # Polygon defining the area of interest
 └── waterbodies.shp   # Polygon labels for water bodies
 ```
@@ -66,6 +67,8 @@ AOI_Name/
 - DTM is used as the reference grid; DSM and intensity will be aligned to it
 - Vector files (shapefiles) can be in any CRS; they will be reprojected as needed
 - NoData values will be handled automatically
+- If provided, `valid_lidar_mask.gpkg` is rasterized to `valid_mask.tif` and used
+  to filter out tiles with insufficient valid coverage (>90% required by default)
 
 ## Usage
 
@@ -147,6 +150,8 @@ python -m geo_deep_learning.train fit --config config/water_extraction_config.ya
      - Computes nDSM = DSM - DTM
      - Computes TWI using WhiteboxTools
      - Stacks layers into multi-band raster: [TWI, nDSM, intensity]
+     - Rasterizes `valid_lidar_mask.gpkg` to `valid_mask.tif` when present and
+       uses it to filter tiles by valid coverage
      - Rasterizes water body labels with AOI masking
      - Tiles into 512×512 patches with 256px stride
    - Generates CSV with train/val/test splits

geo_deep_learning/tools/water_extraction/elevation_stack_datamodule.py

@@ -18,6 +18,7 @@
     compute_twi_whitebox,
     generate_csv_from_tiles,
     rasterize_labels_binary_aoi_mask,
+    rasterize_valid_lidar_mask,
     stack_rasters,
     tile_raster_pair,
 )
@@ -56,6 +57,8 @@ def __init__(  # noqa: PLR0913
         include_intensity: bool = False,
         stride: int = 256,
         test_ratio: float = 0.2,
+        valid_mask_min_ratio: float | None = 0.9,
+        save_rejected_tiles: bool = False,
     ) -> None:
         """
         Initialize ElevationStackDataModule.
@@ -77,6 +80,11 @@ def __init__(  # noqa: PLR0913
                 Defaults to False.
             stride (int): Stride for tiling. Defaults to 256.
             test_ratio (float): Ratio for test split. Defaults to 0.2.
+            valid_mask_min_ratio (float | None): Minimum ratio of valid LiDAR
+                pixels required per tile (uses valid_mask.tif when present).
+                Set to None to disable. Defaults to 0.9.
+            save_rejected_tiles (bool): If True, save tiles filtered out during
+                tiling for debugging. Defaults to False.
 
         """
         super().__init__(
@@ -98,6 +106,8 @@ def __init__(  # noqa: PLR0913
         self.intensity = include_intensity
         self.stride = stride
         self.test_ratio = test_ratio
+        self.valid_mask_min_ratio = valid_mask_min_ratio
+        self.save_rejected_tiles = save_rejected_tiles
 
     def setup(self, stage: str | None = None) -> None:  # noqa: ARG002
         """
@@ -165,6 +175,7 @@ def prepare_data(self) -> None:
             csv_inference_path=self.csv_infer_path,
             test_ratio=self.test_ratio,
             remove_empty_labels=True,
+            valid_mask_min_ratio=self.valid_mask_min_ratio,
         )
 
         # Compute and save statistics
@@ -252,6 +263,7 @@ def _process_aoi(self, aoi_path: str) -> None:
         dsm = Path(aoi_path) / "dsm.tif"
         intensity = Path(aoi_path) / "intensity.tif"
         labels_vector = Path(aoi_path) / "waterbodies.shp"
+        valid_mask_vector = Path(aoi_path) / "valid_lidar_mask.gpkg"
 
         # Step 1: Align inputs to DTM
         log.info("Aligning inputs to DTM")
@@ -299,6 +311,19 @@ def _process_aoi(self, aoi_path: str) -> None:
         log.info("Stacking %d bands: %s", len(stack_inputs), stack_inputs)
         stack_rasters(stack_inputs, str(stack_path))
 
+        # Optional: rasterize valid LiDAR mask
+        valid_mask_raster = out_dir / "valid_mask.tif"
+        if valid_mask_vector.exists():
+            if not valid_mask_raster.exists():
+                log.info("Rasterizing valid LiDAR mask: %s", valid_mask_vector)
+                rasterize_valid_lidar_mask(
+                    str(valid_mask_vector),
+                    str(dtm),
+                    str(valid_mask_raster),
+                )
+            else:
+                log.info("Skipping valid mask rasterization (already exists)")
+
         # Step 4: Rasterize labels
         label_raster = out_dir / "labels_aligned.tif"
         rasterize_labels_binary_aoi_mask(
@@ -319,6 +344,9 @@ def _process_aoi(self, aoi_path: str) -> None:
             output_dir=str(out_dir / "tiles"),
             patch_size=self.patch_size[0],
             stride=self.stride,
+            valid_mask_path=str(valid_mask_raster) if valid_mask_raster.exists() else None,
+            valid_mask_min_ratio=self.valid_mask_min_ratio,
+            save_rejected_tiles=self.save_rejected_tiles,
         )
 
     def train_dataloader(self) -> DataLoader[Any]: