Update video_detection.py

video_detection.py

@@ -53,49 +53,58 @@ def video_detection(cap):
         frame = cv2.resize( frame, (0,0), fx=0.5, fy=0.5)
         #frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
         #if n_frames % SUBSAMPLE == 0:
-        batch.append(frame)
+        #batch.append(frame)
         #if len(batch) == 2 * desired_fps:
-        if len(batch) == 4:
-            #inputs = image_processor(images=batch, return_tensors="pt").to("cuda")
-
-            print(f"starting batch of size {len(batch)}")
-            start = time.time()
-            #with torch.no_grad():
-            #    outputs = model(**inputs)
-            results = model(source=batch, stream=True)
-            end = time.time()
-            print("time taken for inference", end - start)
-
-            start = time.time()
-            #boxes = image_processor.post_process_object_detection(
-            #    outputs,
-             #   target_sizes=torch.tensor([(height, width)] * len(batch)),
-            #    threshold=conf_threshold)
-            """
-            for i, (array, box) in enumerate(zip(batch, boxes)):
-                pil_image = draw_bounding_boxes(Image.fromarray(array), box, model, conf_threshold)
-                frame = np.array(pil_image)
-                # Convert RGB to BGR
-                frame = frame[:, :, ::-1].copy()
-                segment_file.write(frame)
-            """
-            for i, r in enumerate(results):
-                # Plot results image
-                im_bgr = r.plot()  # BGR-order numpy array
-                im_rgb = Image.fromarray(im_bgr[..., ::-1])  # RGB-order PIL image
-                frame = np.array(im_rgb)
-                # Convert RGB to BGR
-                frame = frame[:, :, ::-1].copy()
-                segment_file.write(frame)           
-
-            batch = []
+        #if len(batch) == 4:
+        #inputs = image_processor(images=batch, return_tensors="pt").to("cuda")
+
+        print(f"starting batch of size {len(batch)}")
+        start = time.time()
+        #with torch.no_grad():
+        #    outputs = model(**inputs)
+        results = model(frame, stream=True)
+        end = time.time()
+        print("time taken for inference", end - start)
+
+        start = time.time()
+        #boxes = image_processor.post_process_object_detection(
+        #    outputs,
+         #   target_sizes=torch.tensor([(height, width)] * len(batch)),
+        #    threshold=conf_threshold)
+        """
+        for i, (array, box) in enumerate(zip(batch, boxes)):
+            pil_image = draw_bounding_boxes(Image.fromarray(array), box, model, conf_threshold)
+            frame = np.array(pil_image)
+            # Convert RGB to BGR
+            frame = frame[:, :, ::-1].copy()
+            segment_file.write(frame)
+        """
+        for i, r in enumerate(results):
+            # Plot results image
+            im_bgr = r.plot()  # BGR-order numpy array
+            im_rgb = Image.fromarray(im_bgr[..., ::-1])  # RGB-order PIL image
+            frame = np.array(im_rgb)
+            # Convert RGB to BGR
+            frame = frame[:, :, ::-1].copy()
+            segment_file.write(frame)
+            
+        if n_frames == 3 * fps:
+            n_chunks += 1
             segment_file.release()
-            yield None, name
-            end = time.time()
-            print("time taken for processing boxes", end - start)
+            n_frames = 0
+            yield  frame, name
+            #name = f"output_{n_chunks}{'.mp4' if stream_as_mp4 else '.ts'}"
             name = f"output_{uuid.uuid4()}.mp4"
             segment_file = cv2.VideoWriter(name, video_codec, fps, (width, height)) # type: ignore
 
+        #batch = []
+        #segment_file.release()
+        #yield None, name
+        #end = time.time()
+        #print("time taken for processing boxes", end - start)
+        #name = f"output_{uuid.uuid4()}.mp4"
+        #segment_file = cv2.VideoWriter(name, video_codec, fps, (width, height)) # type: ignore
+
         iterating, frame = cap.read()
         n_frames += 1
 """