117 lines
2.9 KiB
Rust
117 lines
2.9 KiB
Rust
//! 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::BufReader;
|
|
use std::num::NonZeroUsize;
|
|
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,
|
|
#[serde(default)]
|
|
pub auto_save_interval: Option<u32>,
|
|
#[serde(default)]
|
|
pub processing_time: Option<f64>,
|
|
#[serde(default)]
|
|
pub avg_time_per_frame: Option<f64>,
|
|
#[serde(default)]
|
|
pub last_saved_at: Option<String>,
|
|
#[serde(default)]
|
|
pub completed_at: Option<String>,
|
|
#[serde(default)]
|
|
pub auto_save_count: Option<u32>,
|
|
}
|
|
|
|
#[derive(Debug, Clone, Deserialize)]
|
|
pub struct YoloData {
|
|
pub metadata: YoloMetadata,
|
|
pub frames: HashMap<String, FrameData>,
|
|
}
|
|
|
|
pub struct YoloLoader {
|
|
data: YoloData,
|
|
cache: LruCache<u64, Vec<Detection>>,
|
|
}
|
|
|
|
impl YoloLoader {
|
|
const CACHE_SIZE: usize = 60;
|
|
|
|
pub fn load(path: &Path) -> Result<Self> {
|
|
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 cache = LruCache::new(NonZeroUsize::new(Self::CACHE_SIZE).unwrap());
|
|
|
|
Ok(Self { data, cache })
|
|
}
|
|
|
|
pub fn get_detections(&mut self, frame: u64) -> Vec<Detection> {
|
|
if let Some(dets) = self.cache.get(&frame) {
|
|
return dets.clone();
|
|
}
|
|
|
|
if let Some(frame_data) = self.data.frames.get(&frame.to_string()) {
|
|
let dets = frame_data.detections.clone();
|
|
self.cache.put(frame, dets.clone());
|
|
dets
|
|
} 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
|
|
}
|
|
}
|