Added ultralytics

main.py

@@ -3,7 +3,7 @@ import mmpose
 from mmpose.apis import MMPoseInferencer
 
 # Ultralytics
-#from ultralytics import YOLO
+from ultralytics import YOLO
 
 # Gradio
 import gradio as gr
@@ -26,7 +26,7 @@ human3d = MMPoseInferencer(pose3d="human3d")
 # Defining inferencer models to lookup in function
 inferencers = {"Estimate human 2d poses":human, "Estimate human 2d hand poses":hand, "Estimate human 3d poses":human3d}
 
-#track_model = YOLO('yolov8n.pt')  # Load an official Detect model
+track_model = YOLO('yolov8n.pt')  # Load an official Detect model
 
 print("[INFO]: Downloaded models!")
 
@@ -43,49 +43,87 @@ def tracking(video, model, boxes=True):
 
 
 
-def poses(photo, check):
+def show_tracking(video_content, vis_out_dir):
+        video = cv2.VideoCapture(video_content)
+
+        # Track
+        video_track = tracking(video_content, track_model.track)
+
+        # Prepare to save video
+        out_file = os.path.join(vis_out_dir, "track.mp4")
+
+        fourcc = cv2.VideoWriter_fourcc(*"mp4v")  # Codec for MP4 video
+        fps = video.get(cv2.CAP_PROP_FPS)
+        height, width, _ = video_track[0][0].orig_img.shape
+        size = (width,height)
+
+        out_track = cv2.VideoWriter(out_file, fourcc, fps, size)
+
+        # Go through frames and write them 
+        for frame_track in video_track:
+            result_track = frame_track[0].plot()  # plot a BGR numpy array of predictions
+
+        #print(type(result_pose)) numpy ndarray
+        out_track.write(result_track)
+
+        out_track.release()
+
+        video.release()
+        cv2.destroyAllWindows() # Closing window
+
+        return out_file
+
+
+def poses(inferencer, video, vis_out_dir):
+
+    result_generator = inferencer(video, 
+                                vis_out_dir = vis_out_dir,
+                                return_vis=True,
+                                thickness=2,
+                                rebase_keypoint_height=True)    
+    
+    result = [result for result in result_generator] #next(result_generator)    
+
+    out_file = glob.glob(os.path.join(vis_out_dir, "*.mp4"))
+    return out_file
+
+def infer(video, check):
     # Selecting the specific inferencer
     out_files=[]
+
     for i in check:
+         # Create out directory
+        vis_out_dir = str(uuid.uuid4())
         inferencer = inferencers[i] # 'hand', 'human , device='cuda'
 
         print("[INFO]: Running inference!")
-        # Create out directory
-        vis_out_dir = str(uuid.uuid4())
+        out_file = poses(inferencer, video, vis_out_dir)
+        if i == "Detect and track":
+            show_tracking(video, vis_out_dir)
 
-        result_generator = inferencer(photo, 
-                                    vis_out_dir = vis_out_dir,
-                                    return_vis=True,
-                                    thickness=2,
-                                    rebase_keypoint_height=True)    
-        
-        result = [result for result in result_generator] #next(result_generator)    
-
-        out_file = glob.glob(os.path.join(vis_out_dir, "*.mp4"))
-        # 00000.mp4
-        # 000000.mp4
         out_files.append(out_file)
 
     return out_files
 
 def run():
     #https://github.com/open-mmlab/mmpose/blob/main/docs/en/user_guides/inference.md
-    check_web =  gr.CheckboxGroup(choices = ["Estimate human 2d poses", "Estimate human 2d hand poses", "Estimate human 3d poses"], label="Methods", type="value", info="Select the model(s) you want")
-    check_file = gr.CheckboxGroup(choices = ["Estimate human 2d poses", "Estimate human 2d hand poses", "Estimate human 3d poses"], label="Methods", type="value", info="Select the model(s) you want")
+    check_web =  gr.CheckboxGroup(choices = ["Estimate human 2d poses", "Estimate human 2d hand poses", "Estimate human 3d poses", "Detect and track"], label="Methods", type="value", info="Select the model(s) you want")
+    check_file = gr.CheckboxGroup(choices = ["Estimate human 2d poses", "Estimate human 2d hand poses", "Estimate human 3d poses", "Detect and track"], label="Methods", type="value", info="Select the model(s) you want")
+
 
     webcam = gr.Interface(
-        fn=poses,
+        fn=infer,
         inputs= [gr.Video(source="webcam", height=412), check_web],
-        outputs = [gr.PlayableVideo(), gr.PlayableVideo(), gr.PlayableVideo()],
+        outputs = [gr.PlayableVideo(), gr.PlayableVideo(), gr.PlayableVideo(), gr.PlayableVideo()],
         title = 'Pose estimation', 
         description = 'Pose estimation on video',
         allow_flagging=False
         )
 
     file = gr.Interface(
-        poses,
+        infer,
         inputs = [gr.Video(source="upload", height=412), check_file],
-        outputs = [gr.PlayableVideo(),gr.PlayableVideo(),gr.PlayableVideo()],
+        outputs = [gr.PlayableVideo(), gr.PlayableVideo(), gr.PlayableVideo(), gr.PlayableVideo()],
         allow_flagging=False
     )
 
@@ -108,3 +146,7 @@ if __name__ == "__main__":
 # videopose_h36m_27frames_fullconv_supervised-fe8fbba9_20210527.pth
 # videopose_h36m_1frame_fullconv_supervised_cpn_ft-5c3afaed_20210527.pth
 # https://github.com/open-mmlab/mmpose/blob/main/mmpose/apis/inferencers/pose3d_inferencer.py
+
+
+# 00000.mp4
+# 000000.mp4

requirements.txt

@@ -1,3 +1,4 @@
 gradio
 numpy
-opencv-python
+opencv-python
+ultralytics