Update app.py

app.py

@@ -75,6 +75,78 @@ def inference(image_url, image, min_score):
     return out.get_image()
 
 
+def process_vid(video_path):
+
+  cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.3
+  torch.cuda.empty_cache()
+  if not torch.cuda.is_available():
+    cfg.MODEL.DEVICE = "cpu"
+  else:
+    cfg.MODEL.DEVICE = "cuda"
+  predictor = DefaultPredictor(cfg)
+  v = VideoVisualizer(my_metadata,ColorMode.IMAGE)
+  cap = cv2.VideoCapture(video_path)
+  frame_width = int(cap.get(3))
+  frame_height = int(cap.get(4))
+  frame_size = (frame_width,frame_height)
+  fps = int(cap.get(5))
+  vid_fourcc= int(cap.get(cv2.CAP_PROP_FOURCC))
+  output_path = '/content/drive/MyDrive/ColabNotebooks/gradio-exp/output.mp4'
+  fourcc = cv2.VideoWriter_fourcc(*'MJPG')
+  video_writer = cv2.VideoWriter(output_path,fourcc, fps, frame_size)
+
+
+  def runOnVideo(video, maxFrames):
+
+
+    """ Runs the predictor on every frame in the video (unless maxFrames is given),
+    and returns the frame with the predictions drawn.
+    """
+
+    readFrames = 0
+    while True:
+      hasFrame, frame = video.read()
+      if not hasFrame:
+          break
+
+      # Get prediction results for this frame
+      outputs = predictor(frame)
+
+      # Make sure the frame is colored
+      frame = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR)
+
+      # Draw a visualization of the predictions using the video visualizer
+      visualization = v.draw_instance_predictions(frame, outputs["instances"].to("cpu"))
+
+      # Convert Matplotlib RGB format to OpenCV BGR format
+      visualization = cv2.cvtColor(visualization.get_image(), cv2.COLOR_RGB2BGR)
+
+      yield visualization
+
+      readFrames += 1
+      if readFrames > maxFrames:
+          break
+
+    # Create a cut-off for debugging
+  num_frames =  int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
+
+  # Enumerate the frames of the video
+  for visualization in tqdm.tqdm(runOnVideo(cap, num_frames), total=num_frames):
+
+      # Write test image
+      cv2.imwrite('POSE detectron2.png', visualization)
+
+      # Write to video file
+      video_writer.write(visualization)
+
+  # Release resources
+  cap.release()
+  video_writer.release()
+  return output_path
+
+
+
+
 def infer_video(video_path):
     sv.process_video(source_path=video_path, target_path=f"result.mp4", callback=predict_frame)
     return f"result.mp4"
@@ -128,7 +200,7 @@ img_interface = gr.Interface(
    inputs=[input_url,input_image,sliderr], outputs=[output_image], api_name="find"
 )
 video_interface =  gr.Interface(
-    fn=infer_video,
+    fn=process_vid,
    inputs=[input_video], outputs=[output_video], api_name="vid"
 )
 demo = gr.TabbedInterface([img_interface, video_interface], ["Image Upload", "Video Upload"])