f0c876dfc7
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>
280 lines
11 KiB
Python
280 lines
11 KiB
Python
"""Tests for inference functions: depth, edges, segmentation."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from unittest.mock import MagicMock
|
|
from types import SimpleNamespace
|
|
|
|
import numpy as np
|
|
import torch
|
|
|
|
from src.augmentor.inference import (
|
|
compute_edges_from_depth,
|
|
infer_chmv2_batch,
|
|
infer_depth_batch,
|
|
infer_segmentation_batch,
|
|
postprocess_segmentation,
|
|
_SOBEL_X,
|
|
_SOBEL_Y,
|
|
)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# CHMv2 — mock model
|
|
# ---------------------------------------------------------------------------
|
|
|
|
class TestInferChmv2Batch:
|
|
def _make_mock(self, B, H, W):
|
|
model = MagicMock()
|
|
model.parameters = MagicMock(return_value=iter([torch.nn.Parameter(torch.zeros(1))]))
|
|
|
|
def forward(pixel_values=None):
|
|
bb = pixel_values.shape[0]
|
|
return SimpleNamespace(predicted_depth=torch.rand(bb, H, W))
|
|
|
|
model.side_effect = forward
|
|
|
|
processor = MagicMock()
|
|
|
|
def preprocess(images=None, return_tensors=None):
|
|
bb = len(images)
|
|
return {"pixel_values": torch.rand(bb, 3, H, W)}
|
|
|
|
processor.side_effect = preprocess
|
|
|
|
def post_process(outputs, target_sizes=None):
|
|
bb = outputs.predicted_depth.shape[0]
|
|
return [
|
|
{"predicted_depth": torch.rand(target_sizes[i][0], target_sizes[i][1])}
|
|
for i in range(bb)
|
|
]
|
|
|
|
processor.post_process_depth_estimation = post_process
|
|
return model, processor
|
|
|
|
def test_output_shape(self, sample_batch: torch.Tensor) -> None:
|
|
B, _, H, W = sample_batch.shape
|
|
model, processor = self._make_mock(B, H, W)
|
|
result = infer_chmv2_batch(model, processor, sample_batch, torch.device("cpu"))
|
|
assert result.shape == (B, 1, H, W)
|
|
|
|
def test_output_normalized(self, sample_batch: torch.Tensor) -> None:
|
|
B, _, H, W = sample_batch.shape
|
|
model, processor = self._make_mock(B, H, W)
|
|
result = infer_chmv2_batch(model, processor, sample_batch, torch.device("cpu"))
|
|
assert result.min() >= 0.0
|
|
assert result.max() <= 1.0 + 1e-6
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Edges (Sobel) — pure computation, no model needed
|
|
# ---------------------------------------------------------------------------
|
|
|
|
class TestComputeEdges:
|
|
def test_output_shape(self, sample_depth: torch.Tensor) -> None:
|
|
edges = compute_edges_from_depth(sample_depth)
|
|
assert edges.shape == sample_depth.shape
|
|
|
|
def test_output_range(self, sample_depth: torch.Tensor) -> None:
|
|
edges = compute_edges_from_depth(sample_depth)
|
|
assert edges.min() >= 0.0
|
|
assert edges.max() <= 1.0 + 1e-6
|
|
|
|
def test_flat_depth_gives_zero_edges(self) -> None:
|
|
flat = torch.ones(1, 1, 32, 32) * 0.5
|
|
edges = compute_edges_from_depth(flat)
|
|
assert edges.max().item() < 1e-6
|
|
|
|
def test_single_image(self) -> None:
|
|
depth = torch.rand(1, 1, 16, 16)
|
|
edges = compute_edges_from_depth(depth)
|
|
assert edges.shape == (1, 1, 16, 16)
|
|
|
|
def test_batch_consistency(self) -> None:
|
|
"""Batched result matches per-image results."""
|
|
depth = torch.rand(4, 1, 32, 32)
|
|
edges_batch = compute_edges_from_depth(depth)
|
|
for i in range(4):
|
|
edges_single = compute_edges_from_depth(depth[i : i + 1])
|
|
torch.testing.assert_close(edges_batch[i], edges_single[0], atol=1e-5, rtol=1e-5)
|
|
|
|
|
|
class TestSobelKernelCache:
|
|
def test_kernels_are_module_level(self) -> None:
|
|
assert _SOBEL_X.shape == (1, 1, 3, 3)
|
|
assert _SOBEL_Y.shape == (1, 1, 3, 3)
|
|
assert _SOBEL_X.dtype == torch.float32
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Depth — mock model
|
|
# ---------------------------------------------------------------------------
|
|
|
|
class TestInferDepthBatch:
|
|
def test_da3_api_path(self, sample_batch: torch.Tensor) -> None:
|
|
"""Test DA3 branch with model.inference()."""
|
|
B, _, H, W = sample_batch.shape
|
|
|
|
def fake_inference(image_list, process_res=None):
|
|
depth = np.random.rand(len(image_list), H, W).astype(np.float32)
|
|
return SimpleNamespace(depth=depth)
|
|
|
|
model = MagicMock()
|
|
model.inference = fake_inference
|
|
|
|
result = infer_depth_batch(model, sample_batch, torch.device("cpu"))
|
|
assert result.shape == (B, 1, H, W)
|
|
assert result.min() >= 0.0
|
|
assert result.max() <= 1.0 + 1e-6
|
|
|
|
def test_transformers_api_path(self, sample_batch: torch.Tensor) -> None:
|
|
"""Test DA V2 fallback with model(pixel_values=x).predicted_depth."""
|
|
B, _, H, W = sample_batch.shape
|
|
|
|
# Mock a transformers-style model — spec=[] ensures hasattr(model, "inference") is False.
|
|
mock_param = torch.nn.Parameter(torch.zeros(1)) # float32
|
|
model = MagicMock(spec=[])
|
|
model.parameters = MagicMock(return_value=iter([mock_param]))
|
|
|
|
pred_depth = torch.rand(B, H, W)
|
|
model.return_value = SimpleNamespace(predicted_depth=pred_depth)
|
|
|
|
result = infer_depth_batch(model, sample_batch, torch.device("cpu"))
|
|
assert result.shape == (B, 1, H, W)
|
|
assert result.min() >= 0.0
|
|
assert result.max() <= 1.0 + 1e-6
|
|
|
|
def test_per_frame_normalization(self) -> None:
|
|
"""Each frame should be independently normalized to [0, 1]."""
|
|
B, H, W = 2, 16, 16
|
|
|
|
def fake_inference(image_list, process_res=None):
|
|
depth = np.random.rand(len(image_list), H, W).astype(np.float32) * 100
|
|
return SimpleNamespace(depth=depth)
|
|
|
|
model = MagicMock()
|
|
model.inference = fake_inference
|
|
|
|
batch = torch.rand(B, 3, H, W)
|
|
result = infer_depth_batch(model, batch, torch.device("cpu"))
|
|
for i in range(B):
|
|
assert result[i].min().item() < 0.01
|
|
assert result[i].max().item() > 0.99
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Segmentation — mock model
|
|
# ---------------------------------------------------------------------------
|
|
|
|
class TestInferSegmentationBatch:
|
|
def test_segformer_path(self, sample_batch: torch.Tensor) -> None:
|
|
"""Test SegFormer fallback branch."""
|
|
B, _, H, W = sample_batch.shape
|
|
|
|
mock_param = torch.nn.Parameter(torch.zeros(1))
|
|
model = MagicMock()
|
|
model.parameters.return_value = iter([mock_param])
|
|
|
|
num_classes = 5
|
|
logits = torch.randn(B, num_classes, H // 4, W // 4)
|
|
model.return_value = SimpleNamespace(logits=logits)
|
|
|
|
processor = SimpleNamespace(
|
|
image_mean=[0.485, 0.456, 0.406],
|
|
image_std=[0.229, 0.224, 0.225],
|
|
)
|
|
seg_config = {"type": "segformer", "processor": processor}
|
|
|
|
result = infer_segmentation_batch(model, seg_config, sample_batch, torch.device("cpu"))
|
|
assert result.shape == (B, 1, H, W)
|
|
assert result.dtype == torch.uint8
|
|
assert result.min() >= 0
|
|
assert result.max() < num_classes
|
|
|
|
def test_segearth_path(self, sample_batch: torch.Tensor) -> None:
|
|
"""Test SegEarth-OV3 branch with mock pipeline."""
|
|
B, _, H, W = sample_batch.shape
|
|
|
|
def fake_predict(pil_img, text_prompts=None):
|
|
return np.random.randint(0, 3, (H, W), dtype=np.uint8)
|
|
|
|
pipeline = MagicMock()
|
|
pipeline.predict = fake_predict
|
|
|
|
seg_config = {
|
|
"type": "segearth-ov3",
|
|
"prompts": ["bg", "building", "road"],
|
|
}
|
|
|
|
result = infer_segmentation_batch(pipeline, seg_config, sample_batch, torch.device("cpu"))
|
|
assert result.shape == (B, 1, H, W)
|
|
assert result.dtype == torch.uint8
|
|
|
|
def test_segearth_handles_failure(self, sample_batch: torch.Tensor) -> None:
|
|
"""Failed prediction should produce zero tensor, not crash."""
|
|
B, _, H, W = sample_batch.shape
|
|
|
|
pipeline = MagicMock()
|
|
pipeline.predict.side_effect = RuntimeError("OOM")
|
|
|
|
seg_config = {"type": "segearth-ov3", "prompts": ["bg"]}
|
|
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()
|