Update services/video_service.py

services/video_service.py

@@ -1,33 +1,65 @@
 import cv2
-import os
+from PIL import Image
+import numpy as np
+from datetime import datetime
 
-# Global state
-VIDEO_DIR = "data"
-video_files = [os.path.join(VIDEO_DIR, file) for file in sorted(os.listdir(VIDEO_DIR)) if file.endswith((".mp4", ".avi"))]
-video_index = 0
-cap = None
-
-def get_next_video_frame():
-    global cap, video_index
-
-    if not video_files:
-        raise RuntimeError("No video files found in the 'data' directory.")
-
-    if cap is None or not cap.isOpened():
-        cap = cv2.VideoCapture(video_files[video_index])
-
-    ret, frame = cap.read()
-
-    if not ret:
+class VideoService:
+    def load_video(self, video_path):
+        """Load video and return frames."""
+        cap = cv2.VideoCapture(video_path)
+        frames = []
+        while cap.isOpened():
+            ret, frame = cap.read()
+            if not ret:
+                break
+            frames.append(frame)
         cap.release()
-        video_index = (video_index + 1) % len(video_files)
-        cap = cv2.VideoCapture(video_files[video_index])
-        ret, frame = cap.read()
-        if not ret:
-            raise RuntimeError(f"Cannot read video {video_files[video_index]}")
+        return frames
+
+    def frame_to_pil(self, frame):
+        """Convert OpenCV frame to PIL Image."""
+        frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+        return Image.fromarray(frame_rgb)
 
-    return frame
+    def draw_detections(self, frame, detections):
+        """Draw bounding boxes, labels, and timestamp on frame."""
+        frame_copy = frame.copy()
+        # Add timestamp overlay
+        timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
+        cv2.putText(
+            frame_copy,
+            timestamp,
+            (10, 30),
+            cv2.FONT_HERSHEY_SIMPLEX,
+            0.7,
+            (255, 255, 255),
+            2
+        )
+        # Draw bounding boxes and labels for detections
+        for detection in detections:
+            score = detection["score"]
+            label = detection["label"]
+            box = detection["box"]
+            x1, y1, x2, y2 = box["xmin"], box["ymin"], box["xmax"], box["ymax"]
+            cv2.rectangle(frame_copy, (x1, y1), (x2, y2), (0, 255, 0), 2)
+            cv2.putText(
+                frame_copy,
+                f"{label}: {score:.2f}",
+                (x1, y1 - 10),
+                cv2.FONT_HERSHEY_SIMPLEX,
+                0.5,
+                (0, 255, 0),
+                2
+            )
+        return frame_copy
 
-def reset_video_index():
-    global video_index
-    video_index = 0
+    def save_video(self, frames, output_path, fps=30):
+        """Save frames as a video."""
+        if not frames:
+            return
+        height, width, _ = frames[0].shape
+        fourcc = cv2.VideoWriter_fourcc(*"mp4v")
+        out = cv2.VideoWriter(output_path, fourcc, fps, (width, height))
+        for frame in frames:
+            out.write(frame)
+        out.release()