feat(overlay): update YOLO loader
This commit is contained in:
@@ -5,7 +5,8 @@ use lru::LruCache;
|
|||||||
use serde::Deserialize;
|
use serde::Deserialize;
|
||||||
use std::collections::HashMap;
|
use std::collections::HashMap;
|
||||||
use std::fs::File;
|
use std::fs::File;
|
||||||
use std::io::{BufRead, BufReader};
|
use std::io::BufReader;
|
||||||
|
use std::num::NonZeroUsize;
|
||||||
use std::path::Path;
|
use std::path::Path;
|
||||||
|
|
||||||
#[derive(Debug, Clone, Deserialize)]
|
#[derive(Debug, Clone, Deserialize)]
|
||||||
@@ -52,53 +53,39 @@ pub struct YoloData {
|
|||||||
pub struct YoloLoader {
|
pub struct YoloLoader {
|
||||||
data: YoloData,
|
data: YoloData,
|
||||||
cache: LruCache<u64, Vec<Detection>>,
|
cache: LruCache<u64, Vec<Detection>>,
|
||||||
frame_index: HashMap<u64, usize>,
|
|
||||||
file_path: String,
|
|
||||||
}
|
}
|
||||||
|
|
||||||
impl YoloLoader {
|
impl YoloLoader {
|
||||||
const CACHE_SIZE: usize = 60;
|
const CACHE_SIZE: usize = 60;
|
||||||
|
|
||||||
pub fn load(path: &Path) -> Result<Self> {
|
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 file = File::open(path)
|
|
||||||
.with_context(|| format!("Failed to open YOLO file: {:?}", path))?;
|
|
||||||
let reader = BufReader::new(file);
|
let reader = BufReader::new(file);
|
||||||
|
|
||||||
let data: YoloData = serde_json::from_reader(reader)
|
let data: YoloData =
|
||||||
.with_context(|| "Failed to parse YOLO JSON")?;
|
serde_json::from_reader(reader).with_context(|| "Failed to parse YOLO JSON")?;
|
||||||
|
|
||||||
let mut frame_index = HashMap::new();
|
let cache = LruCache::new(NonZeroUsize::new(Self::CACHE_SIZE).unwrap());
|
||||||
for (i, (key, frame)) in data.frames.iter().enumerate() {
|
|
||||||
if let Ok(frame_num) = key.parse::<u64>() {
|
Ok(Self { data, cache })
|
||||||
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> {
|
pub fn get_detections(&mut self, frame: u64) -> Vec<Detection> {
|
||||||
if let Some(dets) = self.cache.get(&frame) {
|
if let Some(dets) = self.cache.get(&frame) {
|
||||||
return dets.iter().collect();
|
return dets.clone();
|
||||||
}
|
}
|
||||||
|
|
||||||
if let Some(frame_data) = self.data.frames.get(&frame.to_string()) {
|
if let Some(frame_data) = self.data.frames.get(&frame.to_string()) {
|
||||||
let dets: Vec<Detection> = frame_data.detections.clone();
|
let dets = frame_data.detections.clone();
|
||||||
self.cache.put(frame, dets.clone());
|
self.cache.put(frame, dets.clone());
|
||||||
dets.iter().collect()
|
dets
|
||||||
} else {
|
} else {
|
||||||
Vec::new()
|
Vec::new()
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
pub fn get_detections_at_time(&mut self, time_ms: u64) -> Vec<&Detection> {
|
pub fn get_detections_at_time(&mut self, time_ms: u64) -> Vec<Detection> {
|
||||||
let fps = self.data.metadata.fps;
|
let fps = self.data.metadata.fps;
|
||||||
let frame = ((time_ms as f64 / 1000.0) * fps) as u64;
|
let frame = ((time_ms as f64 / 1000.0) * fps) as u64;
|
||||||
self.get_detections(frame)
|
self.get_detections(frame)
|
||||||
|
|||||||
Reference in New Issue
Block a user