"""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()