Add YOLO loader module

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<Detection>,
+}
+
+#[derive(Debug, Clone, Deserialize)]
+pub struct YoloMetadata {
+    pub video_path: String,
+    pub model_path: Option<String>,
+    pub width: u32,
+    pub height: u32,
+    pub fps: f64,
+    pub total_frames: u64,
+    pub total_duration: f64,
+    pub processed_at: Option<String>,
+    pub status: Option<String>,
+    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<String, FrameData>,
+}
+
+pub struct YoloLoader {
+    data: YoloData,
+    cache: LruCache<u64, Vec<Detection>>,
+    frame_index: HashMap<u64, usize>,
+    file_path: String,
+}
+
+impl YoloLoader {
+    const CACHE_SIZE: usize = 60;
+
+    pub fn load(path: &Path) -> Result<Self> {
+        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::<u64>() {
+                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<Detection> = 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
+    }
+}