woah

pipeline/sources/file.py

@@ -3,6 +3,7 @@
 import glob
 import logging
 import time
+from pathlib import Path
 from typing import Iterator
 
 import cv2
@@ -19,9 +20,10 @@ class FileStreamSource(StreamSource):
     def __init__(self, path: str):
         self.path = path
         self._fps = 30.0
+        self._image_exts = {".jpg", ".jpeg", ".png", ".bmp", ".webp"}
 
         # Probe for FPS if it's a single file
-        if "*" not in self.path:
+        if "*" not in self.path and not self._is_image_path(self.path):
             try:
                 cap = cv2.VideoCapture(self.path)
                 if cap.isOpened():
@@ -30,6 +32,16 @@ class FileStreamSource(StreamSource):
             except Exception as e:
                 logger.warning("Could not probe FPS from %s: %s", path, e)
 
+    def _is_image_path(self, value: str) -> bool:
+        return Path(value).suffix.lower() in self._image_exts
+
+    def _iter_image_files(self, directory: Path) -> list[str]:
+        files: list[str] = []
+        for ext in self._image_exts:
+            files.extend(glob.glob(str(directory / f"*{ext}")))
+            files.extend(glob.glob(str(directory / f"*{ext.upper()}")))
+        return sorted(set(files))
+
     @property
     def fps(self) -> float:
         return self._fps
@@ -45,6 +57,28 @@ class FileStreamSource(StreamSource):
                     yield Frame(timestamp=time.time(), data=img, frame_id=frame_id)
             return
 
+        file_path = Path(self.path)
+
+        # Handle single image file
+        if self._is_image_path(self.path):
+            img = cv2.imread(self.path)
+            if img is None:
+                logger.warning("Could not read image: %s", self.path)
+                return
+            logger.info("Loaded single image: %s", self.path)
+            yield Frame(timestamp=time.time(), data=img, frame_id=0)
+            return
+
+        # Handle directory of images
+        if file_path.exists() and file_path.is_dir():
+            files = self._iter_image_files(file_path)
+            logger.info("Loading %d images from directory: %s", len(files), self.path)
+            for frame_id, file in enumerate(files):
+                img = cv2.imread(file)
+                if img is not None:
+                    yield Frame(timestamp=time.time(), data=img, frame_id=frame_id)
+            return
+
         logger.info("Opening video file: %s", self.path)
         cap = cv2.VideoCapture(self.path)
         frame_id = 0