diff --git a/src/overlay/yolo.rs b/src/overlay/yolo.rs new file mode 100644 index 0000000..0a6c9a7 --- /dev/null +++ b/src/overlay/yolo.rs @@ -0,0 +1,118 @@ +//! YOLO detection loader with frame caching + +use anyhow::{Context, Result}; +use lru::LruCache; +use serde::Deserialize; +use std::collections::HashMap; +use std::fs::File; +use std::io::{BufRead, BufReader}; +use std::path::Path; + +#[derive(Debug, Clone, Deserialize)] +pub struct Detection { + pub class_id: u32, + pub class_name: String, + pub confidence: f64, + pub x1: f64, + pub y1: f64, + pub x2: f64, + pub y2: f64, +} + +#[derive(Debug, Clone, Deserialize)] +pub struct FrameData { + pub frame_number: u64, + pub time_seconds: f64, + pub time_formatted: String, + pub detections: Vec, +} + +#[derive(Debug, Clone, Deserialize)] +pub struct YoloMetadata { + pub video_path: String, + pub model_path: Option, + pub width: u32, + pub height: u32, + pub fps: f64, + pub total_frames: u64, + pub total_duration: f64, + pub processed_at: Option, + pub status: Option, + pub total_detections: u64, + pub avg_detections_per_frame: f64, +} + +#[derive(Debug, Clone, Deserialize)] +pub struct YoloData { + pub metadata: YoloMetadata, + #[serde(flatten)] + pub frames: HashMap, +} + +pub struct YoloLoader { + data: YoloData, + cache: LruCache>, + frame_index: HashMap, + file_path: String, +} + +impl YoloLoader { + const CACHE_SIZE: usize = 60; + + pub fn load(path: &Path) -> Result { + let file_path = path.to_string_lossy().to_string(); + + let file = File::open(path) + .with_context(|| format!("Failed to open YOLO file: {:?}", path))?; + let reader = BufReader::new(file); + + let data: YoloData = serde_json::from_reader(reader) + .with_context(|| "Failed to parse YOLO JSON")?; + + let mut frame_index = HashMap::new(); + for (i, (key, frame)) in data.frames.iter().enumerate() { + if let Ok(frame_num) = key.parse::() { + frame_index.insert(frame_num, i); + } + } + + Ok(Self { + data, + cache: LruCache::new(Self::CACHE_SIZE), + frame_index, + file_path, + }) + } + + pub fn get_detections(&mut self, frame: u64) -> Vec<&Detection> { + if let Some(dets) = self.cache.get(&frame) { + return dets.iter().collect(); + } + + if let Some(frame_data) = self.data.frames.get(&frame.to_string()) { + let dets: Vec = frame_data.detections.clone(); + self.cache.put(frame, dets.clone()); + dets.iter().collect() + } else { + Vec::new() + } + } + + pub fn get_detections_at_time(&mut self, time_ms: u64) -> Vec<&Detection> { + let fps = self.data.metadata.fps; + let frame = ((time_ms as f64 / 1000.0) * fps) as u64; + self.get_detections(frame) + } + + pub fn metadata(&self) -> &YoloMetadata { + &self.data.metadata + } + + pub fn fps(&self) -> f64 { + self.data.metadata.fps + } + + pub fn total_frames(&self) -> u64 { + self.data.metadata.total_frames + } +}