Initial commit: World-UAV annotation pipeline
4-modality annotation pipeline (depth, edges, segmentation, chmv2) for 973K drone/satellite images. SegEarth-OV3 open-vocabulary segmentation with 11 classes optimized for cross-view geo-localization. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
pikaliov
222
src/tests/test_inference.py
Normal file
222
src/tests/test_inference.py
Normal file
@@ -0,0 +1,222 @@
|
||||
"""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,
|
||||
_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()
|
||||
Reference in New Issue
Block a user