You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
71 lines
2.9 KiB
71 lines
2.9 KiB
"""Streaming CSV detection logger."""
|
|
|
|
import csv
|
|
from pathlib import Path
|
|
from typing import TextIO
|
|
|
|
from .config_loader import DetectionLoggingSettings
|
|
from .models import Detection, FrameInfo
|
|
|
|
|
|
class DetectionLogger:
|
|
"""Write one CSV row per detection, with optional empty-frame rows."""
|
|
|
|
FIELDNAMES = [
|
|
"experiment_id", "frame_number", "timestamp", "elapsed_time_sec", "frame_width", "frame_height",
|
|
"measured_fps", "detection_index", "class_id", "class_name", "confidence", "bbox_x1", "bbox_y1",
|
|
"bbox_x2", "bbox_y2", "bbox_width", "bbox_height", "center_x", "center_y", "foot_x", "foot_y",
|
|
"inside_evaluation_region",
|
|
]
|
|
|
|
def __init__(self, config: DetectionLoggingSettings, experiment_id: str, path: Path) -> None:
|
|
self.config = config
|
|
self.experiment_id = experiment_id
|
|
self.path = path
|
|
self._file: TextIO | None = None
|
|
self._writer: csv.DictWriter[str] | None = None
|
|
self._frames_since_flush = 0
|
|
|
|
def open(self) -> None:
|
|
self.path.parent.mkdir(parents=True, exist_ok=True)
|
|
self._file = self.path.open("w", encoding="utf-8-sig", newline="")
|
|
self._writer = csv.DictWriter(self._file, fieldnames=self.FIELDNAMES)
|
|
self._writer.writeheader()
|
|
|
|
def log_frame(self, info: FrameInfo, detections: list[Detection]) -> None:
|
|
if self._writer is None:
|
|
return
|
|
base = {
|
|
"experiment_id": self.experiment_id, "frame_number": info.frame_number, "timestamp": info.timestamp_iso,
|
|
"elapsed_time_sec": f"{info.elapsed_time_sec:.6f}", "frame_width": info.width, "frame_height": info.height,
|
|
"measured_fps": f"{info.measured_fps:.3f}",
|
|
}
|
|
if detections:
|
|
for index, detection in enumerate(detections):
|
|
b = detection.bbox
|
|
self._writer.writerow(base | {
|
|
"detection_index": index, "class_id": detection.class_id, "class_name": detection.class_name,
|
|
"confidence": f"{detection.confidence:.6f}", "bbox_x1": b.x1, "bbox_y1": b.y1,
|
|
"bbox_x2": b.x2, "bbox_y2": b.y2, "bbox_width": b.width, "bbox_height": b.height,
|
|
"center_x": b.center_x, "center_y": b.center_y, "foot_x": b.foot_x, "foot_y": b.foot_y,
|
|
"inside_evaluation_region": detection.inside_evaluation_region,
|
|
})
|
|
elif self.config.log_empty_frames:
|
|
self._writer.writerow(base)
|
|
self._frames_since_flush += 1
|
|
if self._frames_since_flush >= self.config.flush_interval_frames:
|
|
self._file.flush()
|
|
self._frames_since_flush = 0
|
|
|
|
def close(self) -> None:
|
|
if self._file is not None:
|
|
self._file.close()
|
|
self._file = None
|
|
self._writer = None
|
|
|
|
def __enter__(self) -> "DetectionLogger":
|
|
self.open()
|
|
return self
|
|
|
|
def __exit__(self, *_args: object) -> None:
|
|
self.close()
|
|
|