Add segmentation post-processing: dark water fix + wetland reclassify

Two heuristic rules applied after SegEarth-OV3 inference:

1. Dark water: if background pixels have mean_rgb < 0.24 and std < 0.08,
   reclassify as water. Fixes GTA-UAV satellite dark ocean (57% → ~15% bg).

2. Wetland reclassify (GTA-UAV only): split false-positive wetland pixels
   by color — green-dominant → vegetation, else → bare soil. Fixes 14.3%
   muddy/wetland false positives on GTA-V hillside terrain.

Config flags: seg_fix_dark_water (default True), seg_reclassify_wetland
(default False, enabled in run_gta_uav.py).

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

scripts/run_gta_uav.py

@@ -63,6 +63,8 @@ def main() -> None:
         save_vis=True,
         save_safetensors=True,
         cleanup_npy=True,
+        seg_fix_dark_water=True,
+        seg_reclassify_wetland=True,
         resume=True,
         source=source,
         log_level="INFO",

src/augmentor/inference.py

@@ -300,3 +300,76 @@ def _infer_segformer(
         logits.float(), size=(H, W), mode="bilinear", align_corners=False,
     )
     return upsampled.argmax(dim=1, keepdim=True).cpu().to(torch.uint8)
+
+
+# ---------------------------------------------------------------------------
+# Post-processing heuristics
+# ---------------------------------------------------------------------------
+
+def postprocess_segmentation(
+    seg_ids: torch.Tensor,
+    images_raw: torch.Tensor,
+    water_class: int = 4,
+    dark_water_mean_thr: float = 0.24,
+    dark_water_std_thr: float = 0.08,
+    reclassify_wetland: bool = False,
+    wetland_class: int = 14,
+    vegetation_class: int = 3,
+    bare_soil_class: int = 11,
+    green_thr: float = 0.35,
+) -> torch.Tensor:
+    """Fix known segmentation failure modes with simple heuristics.
+
+    Applied per-image after model inference.
+
+    Rules:
+        1. Dark water: if a background (0) image has mean_rgb < dark_water_mean_thr
+           and std < dark_water_std_thr, reclassify all background pixels as water.
+        2. Wetland reclassification (optional, for GTA-UAV): reclassify wetland pixels
+           by local color — green-dominant → vegetation, else → bare soil.
+
+    Args:
+        seg_ids: [B, 1, H, W] uint8 class IDs.
+        images_raw: [B, 3, H, W] float32 [0, 1] original RGB.
+        water_class: class ID for water.
+        dark_water_mean_thr: mean RGB threshold (0-1) below which bg → water.
+        dark_water_std_thr: std threshold below which bg → water.
+        reclassify_wetland: if True, split wetland into vegetation/bare_soil.
+        wetland_class: class ID for muddy/wetland.
+        vegetation_class: class ID for vegetation.
+        bare_soil_class: class ID for bare soil.
+        green_thr: green-channel ratio threshold for vegetation vs bare soil.
+
+    Returns:
+        seg_ids: [B, 1, H, W] uint8, corrected in-place.
+    """
+    B = seg_ids.shape[0]
+    seg = seg_ids.clone()
+
+    for i in range(B):
+        s = seg[i, 0]        # [H, W] uint8
+        rgb = images_raw[i]  # [3, H, W] float32
+
+        # Rule 1: dark uniform images → background becomes water
+        bg_mask = s == 0
+        if bg_mask.any():
+            bg_pixels = rgb[:, bg_mask]  # [3, N]
+            mean_val = bg_pixels.mean().item()
+            std_val = bg_pixels.std().item()
+            if mean_val < dark_water_mean_thr and std_val < dark_water_std_thr:
+                s[bg_mask] = water_class
+
+        # Rule 2: reclassify wetland by color
+        if reclassify_wetland:
+            wet_mask = s == wetland_class
+            if wet_mask.any():
+                g = rgb[1, wet_mask]  # green channel
+                r = rgb[0, wet_mask]
+                # Green-dominant → vegetation, else → bare soil
+                is_green = g > (r + green_thr * 0.5)
+                reclassed = torch.where(is_green, vegetation_class, bare_soil_class)
+                s[wet_mask] = reclassed.to(torch.uint8)
+
+        seg[i, 0] = s
+
+    return seg

src/conf/pipeline_conf.py

@@ -17,6 +17,8 @@ class PipelineConfig:
         save_concat: bool = False,
         save_safetensors: bool = True,
         cleanup_npy: bool = False,
+        seg_fix_dark_water: bool = True,
+        seg_reclassify_wetland: bool = False,
         resume: bool = True,
         subset: str | None = None,
         source: str | None = None,
@@ -30,6 +32,8 @@ class PipelineConfig:
         self.save_concat = save_concat
         self.save_safetensors = save_safetensors
         self.cleanup_npy = cleanup_npy
+        self.seg_fix_dark_water = seg_fix_dark_water
+        self.seg_reclassify_wetland = seg_reclassify_wetland
         self.resume = resume
         self.subset = subset
         self.source = source

src/main.py

@@ -36,7 +36,7 @@ from src.augmentor.dataset import (
 )
 from src.augmentor.inference import (
     compute_edges_from_depth, infer_depth_batch, infer_chmv2_batch,
-    infer_segmentation_batch,
+    infer_segmentation_batch, postprocess_segmentation,
 )
 from src.augmentor.io_utils import (
     npy_path, vis_path,
@@ -302,6 +302,11 @@ def run_segmentation_stage(
             segs = infer_segmentation_batch(
                 model, seg_config, batch["image_raw"], device,
             )
+            if pipeline_conf.seg_fix_dark_water or pipeline_conf.seg_reclassify_wetland:
+                segs = postprocess_segmentation(
+                    segs, batch["image_raw"],
+                    reclassify_wetland=pipeline_conf.seg_reclassify_wetland,
+                )
             for j in range(segs.shape[0]):
                 save_segmentation_async(
                     segs[j], Path(batch["output_root"][j]),

src/tests/test_inference.py

@@ -13,6 +13,7 @@ from src.augmentor.inference import (
     infer_chmv2_batch,
     infer_depth_batch,
     infer_segmentation_batch,
+    postprocess_segmentation,
     _SOBEL_X,
     _SOBEL_Y,
 )
@@ -220,3 +221,59 @@ class TestInferSegmentationBatch:
         result = infer_segmentation_batch(pipeline, seg_config, sample_batch, torch.device("cpu"))
         assert result.shape == (B, 1, H, W)
         assert (result == 0).all()
+
+
+# ---------------------------------------------------------------------------
+# Post-processing heuristics
+# ---------------------------------------------------------------------------
+
+class TestPostprocessSegmentation:
+    def test_dark_water_reclassified(self) -> None:
+        """Dark uniform background → water."""
+        # Dark image (mean ~0.15, std ~0.02)
+        rgb = torch.full((1, 3, 32, 32), 0.15)
+        rgb += torch.randn_like(rgb) * 0.02
+        rgb.clamp_(0, 1)
+        seg = torch.zeros(1, 1, 32, 32, dtype=torch.uint8)  # all background
+        result = postprocess_segmentation(seg, rgb)
+        assert (result == 4).all()  # water
+
+    def test_bright_background_unchanged(self) -> None:
+        """Bright background should NOT become water."""
+        rgb = torch.full((1, 3, 32, 32), 0.6)
+        seg = torch.zeros(1, 1, 32, 32, dtype=torch.uint8)
+        result = postprocess_segmentation(seg, rgb)
+        assert (result == 0).all()  # stays background
+
+    def test_non_background_unchanged(self) -> None:
+        """Non-background classes untouched even if dark."""
+        rgb = torch.full((1, 3, 32, 32), 0.1)
+        seg = torch.full((1, 1, 32, 32), 3, dtype=torch.uint8)  # vegetation
+        result = postprocess_segmentation(seg, rgb)
+        assert (result == 3).all()
+
+    def test_wetland_reclassify_green(self) -> None:
+        """Green wetland → vegetation when reclassify_wetland=True."""
+        rgb = torch.zeros(1, 3, 32, 32)
+        rgb[0, 1] = 0.6  # green dominant
+        rgb[0, 0] = 0.2  # low red
+        seg = torch.full((1, 1, 32, 32), 14, dtype=torch.uint8)  # wetland
+        result = postprocess_segmentation(seg, rgb, reclassify_wetland=True)
+        assert (result == 3).all()  # vegetation
+
+    def test_wetland_reclassify_brown(self) -> None:
+        """Brown wetland → bare soil when reclassify_wetland=True."""
+        rgb = torch.zeros(1, 3, 32, 32)
+        rgb[0, 0] = 0.5  # red dominant
+        rgb[0, 1] = 0.3  # low green
+        seg = torch.full((1, 1, 32, 32), 14, dtype=torch.uint8)
+        result = postprocess_segmentation(seg, rgb, reclassify_wetland=True)
+        assert (result == 11).all()  # bare soil
+
+    def test_wetland_unchanged_without_flag(self) -> None:
+        """Wetland stays if reclassify_wetland=False."""
+        rgb = torch.zeros(1, 3, 32, 32)
+        rgb[0, 1] = 0.6
+        seg = torch.full((1, 1, 32, 32), 14, dtype=torch.uint8)
+        result = postprocess_segmentation(seg, rgb, reclassify_wetland=False)
+        assert (result == 14).all()