From 27e1434eed9888317489d138f11c967e370aec34 Mon Sep 17 00:00:00 2001
From: Warren Lo <warren@accusys.com.tw>
Date: Thu, 19 Mar 2026 01:26:23 +0800
Subject: [PATCH] feat(overlay): update YOLO loader

---
 src/overlay/yolo.rs | 47 ++++++++++++++++-----------------------------
 1 file changed, 17 insertions(+), 30 deletions(-)

diff --git a/src/overlay/yolo.rs b/src/overlay/yolo.rs
index 0a6c9a7..1186a20 100644
--- a/src/overlay/yolo.rs
+++ b/src/overlay/yolo.rs
@@ -5,7 +5,8 @@ use lru::LruCache;
 use serde::Deserialize;
 use std::collections::HashMap;
 use std::fs::File;
-use std::io::{BufRead, BufReader};
+use std::io::BufReader;
+use std::num::NonZeroUsize;
 use std::path::Path;
 
 #[derive(Debug, Clone, Deserialize)]
@@ -52,53 +53,39 @@ pub struct YoloData {
 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 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,
-        })
+
+        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> {
+    pub fn get_detections(&mut self, frame: u64) -> Vec<Detection> {
         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()) {
-            let dets: Vec<Detection> = frame_data.detections.clone();
+            let dets = frame_data.detections.clone();
             self.cache.put(frame, dets.clone());
-            dets.iter().collect()
+            dets
         } else {
             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 frame = ((time_ms as f64 / 1000.0) * fps) as u64;
         self.get_detections(frame)