rpx_benchmark.tasks.segmentation

 1"""Object (instance) segmentation."""
 2
 3from __future__ import annotations
 4
 5from dataclasses import dataclass
 6
 7from ..api import TaskType
 8from ._pipeline import PipelineResult, TaskRunConfig, run_pipeline
 9from .registry import TaskSpec, register_task
10
11PRIMARY_METRIC = "miou"
12
13
14@dataclass
15class SegmentationRunConfig(TaskRunConfig):
16    pass
17
18
19def run_segmentation(cfg: SegmentationRunConfig) -> PipelineResult:
20    return run_pipeline(
21        task=TaskType.OBJECT_SEGMENTATION,
22        primary_metric=PRIMARY_METRIC,
23        cfg=cfg,
24        compute_ts=True,
25    )
26
27
28def _temporal_stability_hook(predictions, samples, camera_poses):
29    from ..deployment import compute_temporal_stability_seg
30
31    return compute_temporal_stability_seg([p.mask for p in predictions], camera_poses)
32
33
34def _geometric_coherence_hook(predictions, samples):
35    """SGC needs paired mask + depth. Depth arrives via sample.metadata."""
36    import numpy as np
37
38    from ..deployment import compute_sgc
39
40    pairs = [
41        (p.mask, s.metadata["depth_map"])
42        for p, s in zip(predictions, samples, strict=False)
43        if s.metadata and s.metadata.get("depth_map") is not None
44    ]
45    if not pairs:
46        return None
47    masks = [m for m, _ in pairs]
48    depths = [np.asarray(d, dtype=np.float32) for _, d in pairs]
49    return compute_sgc(masks, depths)
50
51
52TASK_SPEC = TaskSpec(
53    task=TaskType.OBJECT_SEGMENTATION,
54    display_name="Object Segmentation",
55    description="Instance mask from a single RGB frame; scored with mIoU.",
56    primary_metric=PRIMARY_METRIC,
57    required_modalities=["rgb", "mask"],
58    higher_is_better=True,
59    run=run_segmentation,
60    temporal_stability_fn=_temporal_stability_hook,
61    geometric_coherence_fn=_geometric_coherence_hook,
62)
63
64register_task(TASK_SPEC)