cuda

main.py

@@ -22,6 +22,12 @@ import uuid
 import numpy as np
 import cv2
 
+# Use GPU if available
+if torch.cuda.is_available():
+    device = torch.device("cuda")
+else:
+    device = torch.device("cpu")
+
 print("[INFO]: Imported modules!")
 human = MMPoseInferencer("human")
 hand = MMPoseInferencer("hand")
@@ -29,11 +35,6 @@ human3d = MMPoseInferencer(pose3d="human3d")
 track_model = YOLO('yolov8n.pt')  # Load an official Detect model
 
 
-if torch.cuda.is_available():
-    device = "cuda"
-else:
-    device = "cpu"
-
 print("[INFO]: Downloaded models!")
 
 def check_extension(video):
@@ -114,11 +115,11 @@ def pose3d(video):
     os.makedirs(vis_out_dir)
 
     result_generator = human3d(video, 
-                                 vis_out_dir = vis_out_dir,
-                                 thickness=2,
-                                 return_vis=True,
-                                 rebase_keypoint_height=True,
-                                 device=device)    
+                                vis_out_dir = vis_out_dir,
+                                thickness=2,
+                                return_vis=True,
+                                rebase_keypoint_height=True,
+                                device=device)    
     
     result = [result for result in result_generator] #next(result_generator)        
 
@@ -231,7 +232,7 @@ def run_UI():
                 
                 gr.Markdown(""" 
                             \n # Information about the models 
-                            
+
                             \n ## Pose models: All the pose estimation models comes from the library [MMpose](https://github.com/open-mmlab/mmpose). It is a library for human pose estimation that provides pre-trained models for 2D and 3D pose estimation. 
 
                             \n ### The 2D pose model is used for estimating the 2D coordinates of human body joints from an image or a video frame. The model uses a convolutional neural network (CNN) to predict the joint locations and their confidence scores.