Update utils/object_detection_brainai.py

utils/object_detection_brainai.py

@@ -39,7 +39,7 @@ class ObjectDetectionModel():
 
         return img_plot, f'Objects Detected: {", ".join(detected_classes) if detected_classes else "No objects detected"}'
 
-    def play_video(self, input_video):
+    def play_video_orig(self, input_video):
         uploaded_video = io.BytesIO(input_video.read())
         temporary_location = "upload.mp4" 
         with open(temporary_location, "wb") as out: 
@@ -75,6 +75,53 @@ class ObjectDetectionModel():
     
         return temp_file.name
 
+    def play_video(self, video_path):
+
+        uploaded_video = io.BytesIO(video_path.read())
+        temporary_location = "upload.mp4"
+        with open(temporary_location, "wb") as temp_out:
+            temp_out.write(uploaded_video.read())
+        temp_out.close()
+        
+        camera = cv2.VideoCapture(temporary_location)
+        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(*'X264')
+        output_video_path = 'output_video.mp4'
+        out = cv2.VideoWriter(output_video_path,fourcc,fps, (frame_width,frame_height))
+        processed_frames = []
+
+        total_frames = int(camera.get(cv2.CAP_PROP_FRAME_COUNT))
+        frame_count = 0
+        progress_bar = st.progress(0)
+        st_frame = st.empty()
+        
+        while(True):
+            ret, frame = camera.read()
+            if not ret:
+                break
+    
+            result = self.model(frame, verbose=False)
+            img_plot = result[0].plot()
+            processed_frames.append(img_plot)
+
+            st_frame.image(img_plot, channels = "BGR")
+            frame_count +=1
+            progress_bar.progress(frame_count/total_frames, text = None)
+            
+        camera.release()
+        for frame in processed_frames:
+             out.write(frame)
+        out.release()
+
+        st_frame.empty()
+        progress_bar.empty()
+
+        return output_video_path
+
+    
     def call_class(self):
         random_class = random.choice(list(self.game_classes))
         return random_class