Update video_detection.py

video_detection.py

@@ -6,6 +6,7 @@ import numpy as np
 import uuid
 
 model = YOLO("model/yolo11n_6-2-25.pt")
+SUBSAMPLE = 2
 
 def draw_boxes(frame, results):
     for r in results:
@@ -31,12 +32,6 @@ def draw_boxes(frame, results):
     return frame
     
 def video_detection(cap):
-    frame_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
-    frame_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
-    fps = int(cap.get(cv2.CAP_PROP_FPS))
-
-    cap = cv2.VideoCapture(video)
-
     video_codec = cv2.VideoWriter_fourcc(*"mp4v") # type: ignore
     fps = int(cap.get(cv2.CAP_PROP_FPS))
 
@@ -51,7 +46,54 @@ def video_detection(cap):
     name = f"output_{uuid.uuid4()}.mp4"
     segment_file = cv2.VideoWriter(name, video_codec, desired_fps, (width, height)) # type: ignore
     batch = []
-
+    while iterating:
+        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)
+        if len(batch) == 2 * desired_fps:
+            #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 = []
+            segment_file.release()
+            yield 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, desired_fps, (width, height)) # type: ignore
+
+        iterating, frame = cap.read()
+        n_frames += 1
 """
 #@spaces.GPU
 def video_detection(cap):