diff --git a/scripts/run_gta_uav.py b/scripts/run_gta_uav.py index f2984f1..a6aab16 100644 --- a/scripts/run_gta_uav.py +++ b/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", diff --git a/src/augmentor/inference.py b/src/augmentor/inference.py index cde5d4b..4869ef1 100644 --- a/src/augmentor/inference.py +++ b/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 diff --git a/src/conf/pipeline_conf.py b/src/conf/pipeline_conf.py index ce199d5..8815ee5 100644 --- a/src/conf/pipeline_conf.py +++ b/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 diff --git a/src/main.py b/src/main.py index f552be1..1636b8b 100644 --- a/src/main.py +++ b/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]), diff --git a/src/tests/test_inference.py b/src/tests/test_inference.py index c19963e..757a493 100644 --- a/src/tests/test_inference.py +++ b/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()