Files
code-test-drone/utils/stereo_ranger/target_store.py
T

120 lines
4.4 KiB
Python

"""目标状态隔离(§4):track_id 优先 + 框位置最近邻兜底。
目标隔离键:track_id 可用(左右一致且 > 0)时优先以 (track, 类别, track_id)
为键——上游 target_tracker 已在链路中,规格 §4 已按此冻结;否则回退类别 +
框位置最近邻(isolation_gate_px 门限,见包 README「实现决策」)。
连续 filter_reset_frames 帧未再出现的目标状态整体清理(位置 + 距离滤波 +
一致性),避免最近邻把新目标接到陈旧状态上。
"""
import math
from dataclasses import dataclass
from typing import Dict, Hashable, Optional, Tuple
from .config import StereoConfig
from .consistency import ConsistencyMonitor
from .distance_filter import DistanceFilter
_KIND_TRACK = "track"
_KIND_NEAREST = "nn"
@dataclass
class _TargetState:
class_id: int
track_id: Optional[int]
center_px: Tuple[float, float]
last_seen_frame: int
class TargetStore:
"""每目标状态容器:隔离键解析 + 帧末老化清理(§11.1 状态清单)。"""
def __init__(self, config: StereoConfig) -> None:
self._config = config
self._distance_filter = DistanceFilter(
config.smooth_alpha, config.jump_reset_m, config.filter_reset_frames
)
self._consistency = ConsistencyMonitor(
config.consistency_jump_px, config.consistency_N_frames, config.consistency_M_frames
)
self._states: Dict[Hashable, _TargetState] = {}
self._next_nearest_id = 0
@property
def distance_filter(self) -> DistanceFilter:
return self._distance_filter
@property
def consistency(self) -> ConsistencyMonitor:
return self._consistency
@property
def count(self) -> int:
"""当前存活的目标状态数。"""
return len(self._states)
def resolve(
self,
class_id: int,
track_id: Optional[int],
center_px: Tuple[float, float],
frame_index: int,
) -> Hashable:
"""解析(或新建)目标 key,并刷新该目标最近位置与出现帧。"""
if track_id is not None and track_id > 0:
key: Hashable = (_KIND_TRACK, class_id, track_id)
self._states[key] = _TargetState(class_id, track_id, center_px, frame_index)
return key
key = self._nearest_key(class_id, center_px, frame_index)
if key is None:
self._next_nearest_id += 1
key = (_KIND_NEAREST, class_id, self._next_nearest_id)
self._states[key] = _TargetState(class_id, None, center_px, frame_index)
else:
state = self._states[key]
state.center_px = center_px
state.last_seen_frame = frame_index
return key
def end_frame(self, frame_index: int) -> None:
"""帧末收尾:未出现的目标记滤波失效,到期整体清理三类状态。"""
stale: List[Hashable] = []
for key, state in self._states.items():
if state.last_seen_frame == frame_index:
continue
self._distance_filter.invalidate(key)
if frame_index - state.last_seen_frame >= self._config.filter_reset_frames:
stale.append(key)
for key in stale:
del self._states[key]
self._distance_filter.drop(key)
self._consistency.drop(key)
def clear(self) -> None:
self._states.clear()
self._distance_filter.clear()
self._consistency.clear()
self._next_nearest_id = 0
def _nearest_key(
self,
class_id: int,
center_px: Tuple[float, float],
frame_index: int,
) -> Optional[Hashable]:
"""同类别、isolation_gate_px 内、本帧尚未被占用的最近状态。"""
gate = self._config.isolation_gate_px
best_key: Optional[Hashable] = None
best_distance: Optional[float] = None
for key, state in self._states.items():
if key[0] != _KIND_NEAREST or state.class_id != class_id:
continue
if state.last_seen_frame == frame_index:
continue # 本帧已被其他框接管,一一归属
distance = math.hypot(center_px[0] - state.center_px[0], center_px[1] - state.center_px[1])
if distance <= gate and (best_distance is None or distance < best_distance):
best_key, best_distance = key, distance
return best_key