Files
NODEDC_MISSION_CORE/src/k1link/perception/detector.py
T

617 lines
23 KiB
Python

"""Versioned fixed-class detector providers for raw-KB4 object proposals."""
from __future__ import annotations
import time
from collections import Counter
from collections.abc import Callable
from dataclasses import dataclass
from threading import Lock
from typing import Final
import numpy as np
from numpy.typing import NDArray
from .contracts import BoundingRegion2D, ObjectProposal2D
from .providers import SourcePacket
from .rf_detr_native_object_detector import (
RF_DETR_NATIVE_CONFIG,
RF_DETR_NATIVE_MODEL_ID,
RF_DETR_NATIVE_MODEL_VERSION,
NativeRfDetrConfig,
NativeRfDetrInferenceBackend,
postprocess_native_rf_detr,
prepare_raw_kb4_rf_detr_native,
)
from .rf_detr_object_detector import (
RF_DETR_CONFIG,
RF_DETR_MODEL_ID,
RF_DETR_MODEL_VERSION,
RfDetrConfig,
RfDetrDetection,
RfDetrInferenceBackend,
postprocess_rf_detr,
preprocess_raw_kb4_rf_detr,
)
from .yolox_object_detector import (
ALL_COCO_YOLOX_CONFIG,
FROZEN_YOLOX_CONFIG,
YOLOX_MODEL_ID,
YOLOX_MODEL_VERSION,
AllCocoYoloxConfig,
FrozenYoloxConfig,
ImageResizer,
InferenceBackend,
YoloxDetection,
postprocess_yolox,
preprocess_raw_kb4,
)
FROZEN_YOLOX_PROVIDER_ID: Final = "triton-yolox-s-raw-kb4/v1"
ALL_COCO_YOLOX_PROVIDER_ID: Final = "triton-yolox-s-raw-kb4-all-coco/v2"
FROZEN_YOLOX_MODEL_ID: Final = f"{YOLOX_MODEL_ID}:{YOLOX_MODEL_VERSION}"
FROZEN_YOLOX_PREPROCESS_ID: Final = "raw-kb4-valid-fov-letterbox/v1"
RF_DETR_SHADOW_PROVIDER_ID: Final = "triton-rf-detr-large-coco-risk-fp16-shadow/v0"
RF_DETR_SHADOW_MODEL_ID: Final = f"{RF_DETR_MODEL_ID}:{RF_DETR_MODEL_VERSION}"
RF_DETR_SHADOW_PREPROCESS_ID: Final = "raw-kb4-valid-fov-rgb-stretch-imagenet/v0"
RF_DETR_NATIVE_SHADOW_PROVIDER_ID: Final = (
"triton-rf-detr-large-coco-native-kb4-risk-fp16-shadow/v0"
)
RF_DETR_NATIVE_SHADOW_MODEL_ID: Final = (
f"{RF_DETR_NATIVE_MODEL_ID}:{RF_DETR_NATIVE_MODEL_VERSION}"
)
RF_DETR_NATIVE_SHADOW_PREPROCESS_ID: Final = (
"raw-kb4-uint8-fused-mask-rgb-pad8-imagenet-trt/v0"
)
class DetectorProviderError(RuntimeError):
"""The detector input, frozen inference or proposal output is incompatible."""
@dataclass(frozen=True, slots=True)
class DetectorProviderSnapshot:
input_frames: int
completed_frames: int
failed_frames: int
zero_proposal_frames: int
proposal_count: int
rejected: tuple[tuple[str, int], ...]
core_duration_ns: int
@dataclass(frozen=True, slots=True)
class DetectorFrameTiming:
sequence: int
preprocess_duration_ns: int
inference_transport_duration_ns: int
postprocess_duration_ns: int
total_duration_ns: int
def __post_init__(self) -> None:
values = (
self.sequence,
self.preprocess_duration_ns,
self.inference_transport_duration_ns,
self.postprocess_duration_ns,
self.total_duration_ns,
)
if any(value < 0 for value in values):
raise DetectorProviderError("detector frame timing must be nonnegative")
if (
self.preprocess_duration_ns
+ self.inference_transport_duration_ns
+ self.postprocess_duration_ns
!= self.total_duration_ns
):
raise DetectorProviderError("detector frame timing does not close")
def to_dict(self) -> dict[str, int]:
return {
"sequence": self.sequence,
"preprocess_duration_ns": self.preprocess_duration_ns,
"inference_transport_duration_ns": self.inference_transport_duration_ns,
"postprocess_duration_ns": self.postprocess_duration_ns,
"total_duration_ns": self.total_duration_ns,
}
DetectorTimingObserver = Callable[[DetectorFrameTiming], None]
@dataclass(frozen=True, slots=True)
class DetectorWarmupSnapshot:
completed: bool
inference_passes: int
preprocess_duration_ns: int
inference_transport_duration_ns: int
postprocess_duration_ns: int
total_duration_ns: int
def __post_init__(self) -> None:
durations = (
self.preprocess_duration_ns,
self.inference_transport_duration_ns,
self.postprocess_duration_ns,
self.total_duration_ns,
)
if (
self.completed is not True
or self.inference_passes != 1
or any(value < 0 for value in durations)
or sum(durations[:3]) != self.total_duration_ns
):
raise DetectorProviderError("detector warmup snapshot is incompatible")
class FrozenYoloxDetectorProvider:
"""One image payload produces one frozen inference request and proposal tuple."""
provider_id: str = FROZEN_YOLOX_PROVIDER_ID
def __init__(
self,
*,
mask: NDArray[np.bool_],
backend: InferenceBackend,
resizer: ImageResizer | None = None,
config: FrozenYoloxConfig | AllCocoYoloxConfig = FROZEN_YOLOX_CONFIG,
clock_ns: Callable[[], int] = time.perf_counter_ns,
) -> None:
if mask.shape != (600, 800) or mask.dtype != np.bool_ or not np.any(mask):
raise DetectorProviderError("frozen valid-FOV mask is incompatible")
self.mask = np.asarray(mask, dtype=np.bool_)
self.backend = backend
self.resizer = resizer
self.config = config
self._clock_ns = clock_ns
self._lock = Lock()
self._input_frames = 0
self._completed_frames = 0
self._failed_frames = 0
self._zero_proposal_frames = 0
self._proposal_count = 0
self._rejected: Counter[str] = Counter()
self._core_duration_ns = 0
def detect(self, packet: SourcePacket) -> tuple[ObjectProposal2D, ...]:
payload = packet.image_payload
with self._lock:
self._input_frames += 1
started_ns = int(self._clock_ns())
try:
if not isinstance(payload, np.ndarray):
raise DetectorProviderError("detector requires a decoded BGR image payload")
image = np.asarray(payload)
if image.dtype != np.uint8:
raise DetectorProviderError("decoded BGR image must be uint8")
tensor = preprocess_raw_kb4(
image,
self.mask,
config=self.config,
resizer=self.resizer,
)
output = self.backend.infer(tensor)
postprocessed = postprocess_yolox(output, self.mask, config=self.config)
proposals = proposals_from_detections(
packet,
postprocessed.detections,
provider_id=self.provider_id,
)
except Exception:
with self._lock:
self._failed_frames += 1
self._core_duration_ns += max(0, int(self._clock_ns()) - started_ns)
raise
with self._lock:
self._completed_frames += 1
self._proposal_count += len(proposals)
self._zero_proposal_frames += not proposals
self._rejected.update(dict(postprocessed.rejected))
self._core_duration_ns += max(0, int(self._clock_ns()) - started_ns)
return proposals
def snapshot(self) -> DetectorProviderSnapshot:
with self._lock:
return DetectorProviderSnapshot(
input_frames=self._input_frames,
completed_frames=self._completed_frames,
failed_frames=self._failed_frames,
zero_proposal_frames=self._zero_proposal_frames,
proposal_count=self._proposal_count,
rejected=tuple(sorted(self._rejected.items())),
core_duration_ns=self._core_duration_ns,
)
def proposals_from_detections(
packet: SourcePacket,
detections: tuple[YoloxDetection, ...],
*,
provider_id: str = FROZEN_YOLOX_PROVIDER_ID,
) -> tuple[ObjectProposal2D, ...]:
envelope = packet.envelope
return tuple(
ObjectProposal2D(
proposal_id=f"proposal-{envelope.sequence}-{index}",
source_id=envelope.source_id,
frame_id=envelope.frame_id,
region=BoundingRegion2D(*detection.bbox_xyxy),
objectness=detection.score,
provider_id=provider_id,
model_id=FROZEN_YOLOX_MODEL_ID,
preprocess_id=FROZEN_YOLOX_PREPROCESS_ID,
semantic_hint=detection.label,
provider_tracklet=None,
)
for index, detection in enumerate(detections)
)
class AllCocoYoloxDetectorProvider(FrozenYoloxDetectorProvider):
"""Emit every qualified COCO class without adding another inference pass."""
provider_id: str = ALL_COCO_YOLOX_PROVIDER_ID
def __init__(
self,
*,
mask: NDArray[np.bool_],
backend: InferenceBackend,
resizer: ImageResizer | None = None,
config: AllCocoYoloxConfig = ALL_COCO_YOLOX_CONFIG,
clock_ns: Callable[[], int] = time.perf_counter_ns,
) -> None:
super().__init__(
mask=mask,
backend=backend,
resizer=resizer,
config=config,
clock_ns=clock_ns,
)
class RfDetrShadowDetectorProvider:
"""Emit behavior-relevant fixed classes from one RF-DETR inference pass."""
provider_id: str = RF_DETR_SHADOW_PROVIDER_ID
def __init__(
self,
*,
mask: NDArray[np.bool_],
backend: RfDetrInferenceBackend,
resizer: ImageResizer | None = None,
config: RfDetrConfig = RF_DETR_CONFIG,
clock_ns: Callable[[], int] = time.perf_counter_ns,
timing_observer: DetectorTimingObserver | None = None,
) -> None:
if mask.shape != (600, 800) or mask.dtype != np.bool_ or not np.any(mask):
raise DetectorProviderError("RF-DETR valid-FOV mask is incompatible")
self.mask = np.asarray(mask, dtype=np.bool_)
self.backend = backend
self.resizer = resizer
self.config = config
self._clock_ns = clock_ns
self.timing_observer = timing_observer
self._lock = Lock()
self._input_frames = 0
self._completed_frames = 0
self._failed_frames = 0
self._zero_proposal_frames = 0
self._proposal_count = 0
self._rejected: Counter[str] = Counter()
self._core_duration_ns = 0
self._warmup_started = False
self._warmup_snapshot: DetectorWarmupSnapshot | None = None
def warm_up(self) -> DetectorWarmupSnapshot:
"""Prime preprocessing, transport and postprocessing before source admission."""
with self._lock:
if self._warmup_snapshot is not None:
return self._warmup_snapshot
if self._warmup_started:
raise DetectorProviderError("RF-DETR warmup is already in progress")
self._warmup_started = True
started_ns = int(self._clock_ns())
try:
image = np.zeros(
(self.config.source_height, self.config.source_width, 3),
dtype=np.uint8,
)
tensor = preprocess_raw_kb4_rf_detr(
image,
self.mask,
config=self.config,
resizer=self.resizer,
)
preprocessed_ns = int(self._clock_ns())
output = self.backend.infer(tensor)
inferred_ns = int(self._clock_ns())
postprocess_rf_detr(output, self.mask, config=self.config)
completed_ns = int(self._clock_ns())
except Exception:
with self._lock:
self._warmup_started = False
raise
snapshot = DetectorWarmupSnapshot(
completed=True,
inference_passes=1,
preprocess_duration_ns=max(0, preprocessed_ns - started_ns),
inference_transport_duration_ns=max(0, inferred_ns - preprocessed_ns),
postprocess_duration_ns=max(0, completed_ns - inferred_ns),
total_duration_ns=max(0, completed_ns - started_ns),
)
with self._lock:
self._warmup_snapshot = snapshot
return snapshot
def detect(self, packet: SourcePacket) -> tuple[ObjectProposal2D, ...]:
payload = packet.image_payload
with self._lock:
self._input_frames += 1
started_ns = int(self._clock_ns())
try:
if not isinstance(payload, np.ndarray):
raise DetectorProviderError("RF-DETR requires a decoded BGR image payload")
image = np.asarray(payload)
if image.dtype != np.uint8:
raise DetectorProviderError("decoded BGR image must be uint8")
tensor = preprocess_raw_kb4_rf_detr(
image,
self.mask,
config=self.config,
resizer=self.resizer,
)
preprocessed_ns = (
int(self._clock_ns()) if self.timing_observer is not None else started_ns
)
output = self.backend.infer(tensor)
inferred_ns = (
int(self._clock_ns()) if self.timing_observer is not None else preprocessed_ns
)
postprocessed = postprocess_rf_detr(output, self.mask, config=self.config)
proposals = proposals_from_rf_detr_detections(packet, postprocessed.detections)
except Exception:
with self._lock:
self._failed_frames += 1
self._core_duration_ns += max(0, int(self._clock_ns()) - started_ns)
raise
completed_ns = int(self._clock_ns())
with self._lock:
self._completed_frames += 1
self._proposal_count += len(proposals)
self._zero_proposal_frames += not proposals
self._rejected.update(dict(postprocessed.rejected))
self._core_duration_ns += max(0, completed_ns - started_ns)
if self.timing_observer is not None:
self.timing_observer(
DetectorFrameTiming(
sequence=packet.envelope.sequence,
preprocess_duration_ns=max(0, preprocessed_ns - started_ns),
inference_transport_duration_ns=max(0, inferred_ns - preprocessed_ns),
postprocess_duration_ns=max(0, completed_ns - inferred_ns),
total_duration_ns=max(0, completed_ns - started_ns),
)
)
return proposals
def snapshot(self) -> DetectorProviderSnapshot:
with self._lock:
return DetectorProviderSnapshot(
input_frames=self._input_frames,
completed_frames=self._completed_frames,
failed_frames=self._failed_frames,
zero_proposal_frames=self._zero_proposal_frames,
proposal_count=self._proposal_count,
rejected=tuple(sorted(self._rejected.items())),
core_duration_ns=self._core_duration_ns,
)
def proposals_from_rf_detr_detections(
packet: SourcePacket,
detections: tuple[RfDetrDetection, ...],
) -> tuple[ObjectProposal2D, ...]:
envelope = packet.envelope
return tuple(
ObjectProposal2D(
proposal_id=f"proposal-{envelope.sequence}-{index}",
source_id=envelope.source_id,
frame_id=envelope.frame_id,
region=BoundingRegion2D(*detection.bbox_xyxy),
objectness=detection.score,
provider_id=RF_DETR_SHADOW_PROVIDER_ID,
model_id=RF_DETR_SHADOW_MODEL_ID,
preprocess_id=RF_DETR_SHADOW_PREPROCESS_ID,
semantic_hint=detection.label,
provider_tracklet=None,
)
for index, detection in enumerate(detections)
)
class NativeRfDetrShadowDetectorProvider:
"""Emit risk classes from one exact-raster native RF-DETR inference pass."""
provider_id: str = RF_DETR_NATIVE_SHADOW_PROVIDER_ID
def __init__(
self,
*,
mask: NDArray[np.bool_],
backend: NativeRfDetrInferenceBackend,
config: NativeRfDetrConfig = RF_DETR_NATIVE_CONFIG,
clock_ns: Callable[[], int] = time.perf_counter_ns,
timing_observer: DetectorTimingObserver | None = None,
) -> None:
if mask.shape != (600, 800) or mask.dtype != np.bool_ or not np.any(mask):
raise DetectorProviderError("native RF-DETR valid-FOV mask is incompatible")
self.mask = np.asarray(mask, dtype=np.bool_)
self.backend = backend
self.config = config
self._clock_ns = clock_ns
self.timing_observer = timing_observer
self._lock = Lock()
self._input_frames = 0
self._completed_frames = 0
self._failed_frames = 0
self._zero_proposal_frames = 0
self._proposal_count = 0
self._rejected: Counter[str] = Counter()
self._core_duration_ns = 0
self._warmup_started = False
self._warmup_snapshot: DetectorWarmupSnapshot | None = None
def warm_up(self) -> DetectorWarmupSnapshot:
"""Prime raw transport and postprocessing before source admission."""
with self._lock:
if self._warmup_snapshot is not None:
return self._warmup_snapshot
if self._warmup_started:
raise DetectorProviderError("native RF-DETR warmup is already in progress")
self._warmup_started = True
started_ns = int(self._clock_ns())
try:
image = np.zeros(
(self.config.source_height, self.config.source_width, 3),
dtype=np.uint8,
)
tensor = prepare_raw_kb4_rf_detr_native(image, config=self.config)
preprocessed_ns = int(self._clock_ns())
output = self.backend.infer(tensor)
inferred_ns = int(self._clock_ns())
postprocess_native_rf_detr(output, self.mask, config=self.config)
completed_ns = int(self._clock_ns())
except Exception:
with self._lock:
self._warmup_started = False
raise
snapshot = DetectorWarmupSnapshot(
completed=True,
inference_passes=1,
preprocess_duration_ns=max(0, preprocessed_ns - started_ns),
inference_transport_duration_ns=max(0, inferred_ns - preprocessed_ns),
postprocess_duration_ns=max(0, completed_ns - inferred_ns),
total_duration_ns=max(0, completed_ns - started_ns),
)
with self._lock:
self._warmup_snapshot = snapshot
return snapshot
def detect(self, packet: SourcePacket) -> tuple[ObjectProposal2D, ...]:
payload = packet.image_payload
with self._lock:
self._input_frames += 1
started_ns = int(self._clock_ns())
try:
if not isinstance(payload, np.ndarray):
raise DetectorProviderError(
"native RF-DETR requires a decoded BGR image payload"
)
image = np.asarray(payload)
if image.dtype != np.uint8:
raise DetectorProviderError("decoded BGR image must be uint8")
tensor = prepare_raw_kb4_rf_detr_native(image, config=self.config)
preprocessed_ns = (
int(self._clock_ns()) if self.timing_observer is not None else started_ns
)
output = self.backend.infer(tensor)
inferred_ns = (
int(self._clock_ns()) if self.timing_observer is not None else preprocessed_ns
)
postprocessed = postprocess_native_rf_detr(
output,
self.mask,
config=self.config,
)
proposals = proposals_from_native_rf_detr_detections(
packet,
postprocessed.detections,
)
except Exception:
with self._lock:
self._failed_frames += 1
self._core_duration_ns += max(0, int(self._clock_ns()) - started_ns)
raise
completed_ns = int(self._clock_ns())
with self._lock:
self._completed_frames += 1
self._proposal_count += len(proposals)
self._zero_proposal_frames += not proposals
self._rejected.update(dict(postprocessed.rejected))
self._core_duration_ns += max(0, completed_ns - started_ns)
if self.timing_observer is not None:
self.timing_observer(
DetectorFrameTiming(
sequence=packet.envelope.sequence,
preprocess_duration_ns=max(0, preprocessed_ns - started_ns),
inference_transport_duration_ns=max(0, inferred_ns - preprocessed_ns),
postprocess_duration_ns=max(0, completed_ns - inferred_ns),
total_duration_ns=max(0, completed_ns - started_ns),
)
)
return proposals
def snapshot(self) -> DetectorProviderSnapshot:
with self._lock:
return DetectorProviderSnapshot(
input_frames=self._input_frames,
completed_frames=self._completed_frames,
failed_frames=self._failed_frames,
zero_proposal_frames=self._zero_proposal_frames,
proposal_count=self._proposal_count,
rejected=tuple(sorted(self._rejected.items())),
core_duration_ns=self._core_duration_ns,
)
def proposals_from_native_rf_detr_detections(
packet: SourcePacket,
detections: tuple[RfDetrDetection, ...],
) -> tuple[ObjectProposal2D, ...]:
envelope = packet.envelope
return tuple(
ObjectProposal2D(
proposal_id=f"proposal-{envelope.sequence}-{index}",
source_id=envelope.source_id,
frame_id=envelope.frame_id,
region=BoundingRegion2D(*detection.bbox_xyxy),
objectness=detection.score,
provider_id=RF_DETR_NATIVE_SHADOW_PROVIDER_ID,
model_id=RF_DETR_NATIVE_SHADOW_MODEL_ID,
preprocess_id=RF_DETR_NATIVE_SHADOW_PREPROCESS_ID,
semantic_hint=detection.label,
provider_tracklet=None,
)
for index, detection in enumerate(detections)
)
__all__ = [
"ALL_COCO_YOLOX_PROVIDER_ID",
"FROZEN_YOLOX_MODEL_ID",
"FROZEN_YOLOX_PREPROCESS_ID",
"FROZEN_YOLOX_PROVIDER_ID",
"RF_DETR_SHADOW_MODEL_ID",
"RF_DETR_SHADOW_PREPROCESS_ID",
"RF_DETR_SHADOW_PROVIDER_ID",
"RF_DETR_NATIVE_SHADOW_MODEL_ID",
"RF_DETR_NATIVE_SHADOW_PREPROCESS_ID",
"RF_DETR_NATIVE_SHADOW_PROVIDER_ID",
"DetectorProviderError",
"DetectorProviderSnapshot",
"DetectorFrameTiming",
"DetectorTimingObserver",
"DetectorWarmupSnapshot",
"AllCocoYoloxDetectorProvider",
"FrozenYoloxDetectorProvider",
"NativeRfDetrShadowDetectorProvider",
"RfDetrShadowDetectorProvider",
"proposals_from_detections",
"proposals_from_native_rf_detr_detections",
"proposals_from_rf_detr_detections",
]