Update utils/object_detection_brainai.py

utils/object_detection_brainai.py

@@ -1,11 +1,9 @@
-from ultralytics import YOLO
 import cv2
+from ultralytics import YOLO
 
-class ObjectDetectionModel(): 
-    
+class ObjectDetectionModel():
     def __init__(self):
         self.model = YOLO("models/yolov8n_openvino_model", task = "detect")
-        self.class_names = self.model.names
 
     def process(self, img):
         result = self.model(img)
@@ -13,33 +11,29 @@ class ObjectDetectionModel():
 
         return img_plot
 
-    
-    def detect_objects_in_video(self, video_path, conf_thres=0.25, iou_thres=0.45):
-        cap = cv2.VideoCapture(video_path)
-        frame_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
-        frame_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
-        fps = cap.get(cv2.CAP_PROP_FPS)
-        frames = []
-    
-        while cap.isOpened():
-            ret, frame = cap.read()
+    def play_video(self, video_path):
+        camera = cv2.VideoCapture(video_path)
+        frame_width = int(camera.get(cv2.CAP_PROP_FRAME_WIDTH))
+        frame_height = int(camera.get(cv2.CAP_PROP_FRAME_HEIGHT))
+        fps = camera.get(cv2.CAP_PROP_FPS)
+
+        fourcc = cv2.VideoWriter_fourcc(*'mp4v')
+        output_video_path = 'output_video.mp4'
+        out = cv2.VideoWriter(output_video_path,fourcc,fps, 
+                              (frame_width,frame_height))
+        processed_frames = []
+        
+        while(True):
+            ret, frame = camera.read()
             if not ret:
                 break
     
-            img_plot = self.process(frame)
-            frames.append(img_plot)
-        cap.release()
-        return frames, fps, frame_width, frame_height
-    
-    def create_video(self, frames, fps, frame_width, frame_height):
-        fourcc = cv2.VideoWriter_fourcc(*'mp4v')
-        out_path = 'output_video.mp4'
-        out = cv2.VideoWriter(out_path, fourcc, fps, (frame_width, frame_height))
-    
-        for frame in frames:
-            out.write(frame)
-    
+            result = self.model(frame, verbose=False)
+            img_plot = result[0].plot()
+            processed_frames.append(img_plot)
+        camera.release()
+        for frame in processed_frames:
+             out.write(frame)
         out.release()
-        return out_path
 
-    
+        return output_video_path