cuda

Dockerfile

@@ -4,7 +4,8 @@ WORKDIR /code
 COPY ./requirements.txt /code/requirements.txt
 
 #RUN pip install torch==1.9.1+cu111 torchvision==0.10.1+cu111 -f https://download.pytorch.org/whl/torch_stable.ht
-RUN pip install torch==1.9.1 torchvision==0.10.1 -f https://download.pytorch.org/whl/torch_stable.ht
+RUN pip install torch==1.9.1 torchvision==0.10.1 -f https://download.pytorch.org/whl/cu113 
+#https://download.pytorch.org/whl/torch_stable.ht
 RUN pip install --no-cache-dir --upgrade -U openmim
 RUN mim install --no-cache-dir --upgrade mmengine
 RUN mim install "mmcv>=2.0.1"

main.py

@@ -29,6 +29,8 @@ human3d = MMPoseInferencer(pose3d="human3d")
 track_model = YOLO('yolov8n.pt')  # Load an official Detect model
 
 # ultraltics
+if torch.cuda.is_available():
+    device = "cuda"
 
 print("[INFO]: Downloaded models!")
 
@@ -38,7 +40,6 @@ def check_extension(video):
     # extract the file name and extension
     file_name = split_tup[0]
     file_extension = split_tup[1]
-    print(file_extension)
 
     if file_extension != ".mp4":
         print("Converting to mp4")
@@ -50,9 +51,6 @@ def check_extension(video):
     return video
 
 
-
-
-
 def tracking(video, model, boxes=True):
     print("[INFO] Is cuda available? ", torch.cuda.is_available())
     print("[INFO] Loading model...")
@@ -61,12 +59,11 @@ def tracking(video, model, boxes=True):
     # Perform tracking with the model
     print("[INFO] Starting tracking!")
     # https://docs.ultralytics.com/modes/predict/
-    annotated_frame = model(video, boxes=boxes)
+    annotated_frame = model(video, boxes=boxes, device=device)
 
     return annotated_frame
 
 def show_tracking(video_content):
-        print()
 
         # https://docs.ultralytics.com/datasets/detect/coco/
         video = cv2.VideoCapture(video_content)
@@ -116,7 +113,7 @@ def pose3d(video):
                                  thickness=2,
                                  return_vis=True,
                                  rebase_keypoint_height=True,
-                                 device="cuda")    
+                                 device=device)    
     
     result = [result for result in result_generator] #next(result_generator)        
 
@@ -139,7 +136,7 @@ def pose2d(video, kpt_threshold):
                             thickness=2,
                             rebase_keypoint_height=True,
                             kpt_thr=kpt_threshold,
-                            device="cuda"
+                            device=device
                             )    
     
     result = [result for result in result_generator] #next(result_generator)        
@@ -163,7 +160,7 @@ def pose2dhand(video, kpt_threshold):
                                  thickness=2,
                                  rebase_keypoint_height=True,
                                  kpt_thr=kpt_threshold,
-                                 device="cuda")    
+                                 device=device)    
     
     result = [result for result in result_generator] #next(result_generator)