Create video_utils.py

video_utils.py

@@ -0,0 +1,49 @@
+import cv2
+import os
+from PIL import Image
+from model import predict_fault
+
+def extract_frames(video_path, interval=30):
+    cap = cv2.VideoCapture(video_path)
+    frames = []
+    frame_count = 0
+
+    while cap.isOpened():
+        ret, frame = cap.read()
+        if not ret:
+            break
+
+        if frame_count % interval == 0:
+            img = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+            pil_img = Image.fromarray(img)
+            frames.append((frame_count, pil_img))
+        
+        frame_count += 1
+
+    cap.release()
+    return frames
+
+def annotate_video(video_path, predictions: dict, output_path="output.mp4"):
+    cap = cv2.VideoCapture(video_path)
+    width  = int(cap.get(3))
+    height = int(cap.get(4))
+    fps    = cap.get(cv2.CAP_PROP_FPS)
+
+    out = cv2.VideoWriter(output_path, cv2.VideoWriter_fourcc(*'mp4v'), fps, (width, height))
+    frame_index = 0
+
+    while cap.isOpened():
+        ret, frame = cap.read()
+        if not ret:
+            break
+
+        if frame_index in predictions:
+            label = predictions[frame_index]
+            cv2.putText(frame, label, (10, 40), cv2.FONT_HERSHEY_SIMPLEX, 1.2, (0, 0, 255), 2)
+
+        out.write(frame)
+        frame_index += 1
+
+    cap.release()
+    out.release()
+    return output_path