Update utils/object_detection_brainai.py

utils/object_detection_brainai.py

@@ -16,7 +16,6 @@ class ObjectDetectionModel():
 
         with open("utils/game_classes.txt", "r") as file:
            self.game_classes = [line.strip() for line in file if line.strip()]
-
     
     def process_image(self, img):
 
@@ -39,37 +38,38 @@ class ObjectDetectionModel():
 
         return img_plot, f'Objects Detected: {", ".join(detected_classes) if detected_classes else "No objects detected"}'
 
-    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()
+    def play_video(self, input_video):
+        uploaded_video = io.BytesIO(input_video.read())
+        temporary_location = "upload.mp4" 
+        with open(temporary_location, "wb") as out: 
+            out.write(uploaded_video.read())
+        out.close() 
         
         camera = cv2.VideoCapture(temporary_location)
         fps = camera.get(cv2.CAP_PROP_FPS)
         temp_file = tempfile.NamedTemporaryFile(delete=False, suffix='.mp4')
-        processed_frames = []
-
+        processed_frames=[]
         total_frames = int(camera.get(cv2.CAP_PROP_FRAME_COUNT))
-        frame_count = 0
         progress_bar = st.progress(0)
+        frame_count = 0
         st_frame = st.empty()
         
-        while(True):
+        while(camera.isOpened()):
             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)
-            
+            if ret:
+                img_plot, _ = self.process_image(frame)
+                st_frame.image(img_plot, channels = "BGR")
+                processed_frames.append(cv2.cvtColor(img_plot,cv2.COLOR_BGR2RGB))
+                frame_count +=1
+                progress_bar.progress(frame_count/total_frames, text = None)
+    
+            else:
+                camera.release()
+                st_frame.empty()
+                progress_bar.empty()
+                break
+                
         camera.release()
         clip = mpy.ImageSequenceClip(processed_frames, fps=fps)
         clip.write_videofile(temp_file.name)
@@ -78,8 +78,8 @@ class ObjectDetectionModel():
         progress_bar.empty()
 
         return temp_file.name
-    
+
     def call_class(self):
         random_class = random.choice(list(self.game_classes))
         return random_class
-    
+