diff --git a/croppa/main.py b/croppa/main.py index 03b092c..a091d0f 100644 --- a/croppa/main.py +++ b/croppa/main.py @@ -12,6 +12,7 @@ import json import subprocess import queue import ctypes +from collections import OrderedDict from PIL import Image def load_image_utf8(image_path): @@ -279,7 +280,7 @@ class FeatureTracker: class Cv2BufferedCap: """Buffered wrapper around cv2.VideoCapture that handles frame loading, seeking, and caching correctly""" - def __init__(self, video_path, backend=None): + def __init__(self, video_path, backend=None, cache_size=10000): self.video_path = video_path self.cap = cv2.VideoCapture(str(video_path), backend) if not self.cap.isOpened(): @@ -294,6 +295,10 @@ class Cv2BufferedCap: # Current position tracking self.current_frame = 0 + # Frame cache (LRU) + self.cache_size = cache_size + self.frame_cache = OrderedDict() + def get_frame(self, frame_number): @@ -301,6 +306,11 @@ class Cv2BufferedCap: # Clamp frame number to valid range frame_number = max(0, min(frame_number, self.total_frames - 1)) + # Check cache first + if frame_number in self.frame_cache: + self.frame_cache.move_to_end(frame_number) + return self.frame_cache[frame_number] + # Optimize for sequential reading (next frame) if frame_number == self.current_frame + 1: ret, frame = self.cap.read() @@ -311,6 +321,11 @@ class Cv2BufferedCap: if ret: self.current_frame = frame_number + # Store in cache, evict least recently used if cache is full + if len(self.frame_cache) >= self.cache_size: + self.frame_cache.popitem(last=False) + self.frame_cache[frame_number] = frame + self.frame_cache.move_to_end(frame_number) return frame else: raise ValueError(f"Failed to read frame {frame_number}")