diff --git a/mmpose/loop_mmpose.py b/mmpose/loop_mmpose.py index 6f18a880..eed0e461 100644 --- a/mmpose/loop_mmpose.py +++ b/mmpose/loop_mmpose.py @@ -35,7 +35,13 @@ def checkCudaPyTorch(): if os.path.isfile(video_path): os.remove(video_path) -checkCudaPyTorch() +# checkCudaPyTorch() +if torch.cuda.is_available(): + num_gpus = torch.cuda.device_count() + logging.info(f"Found {num_gpus} GPU(s).") +else: + logging.info("No GPU detected. Exiting.") + raise Exception("No GPU detected. Exiting.") while True: if not os.path.isfile(video_path): time.sleep(0.1) @@ -48,7 +54,13 @@ def checkCudaPyTorch(): os.makedirs(output_dir) try: - checkCudaPyTorch() + # checkCudaPyTorch() + if torch.cuda.is_available(): + num_gpus = torch.cuda.device_count() + logging.info(f"Found {num_gpus} GPU(s).") + else: + logging.info("No GPU detected. Exiting.") + raise Exception("No GPU detected. Exiting.") # Run human detection. pathModelCkptPerson = model_ckpt_person bboxPath = os.path.join(output_dir, 'box.pkl') diff --git a/openpose/loop_openpose.py b/openpose/loop_openpose.py index 77acbb10..87d0a9a3 100644 --- a/openpose/loop_openpose.py +++ b/openpose/loop_openpose.py @@ -75,7 +75,16 @@ def getResolutionCommand(resolutionPoseDetection, horizontal): #%% logging.info("Waiting for data...") -check_cuda_device() +# check_cuda_device() +try: + # Run the nvidia-smi command and capture the output + _ = subprocess.check_output(["nvidia-smi"]) + # If the command ran successfully, assume a CUDA device is present + logging.info("A CUDA-capable device is detected.") +except subprocess.CalledProcessError as e: + # If the command fails, it means no CUDA device is detected + logging.info("No CUDA-capable device is detected. Error:", e) + raise Exception("No CUDA-capable device is detected.") video_path = "/openpose/data/video_openpose.mov" output_dir = "/openpose/data/output_openpose" @@ -102,7 +111,16 @@ def getResolutionCommand(resolutionPoseDetection, horizontal): horizontal = getVideoOrientation(video_path) cmd_hr = getResolutionCommand(resolutionPoseDetection, horizontal) - check_cuda_device() + # check_cuda_device() + try: + # Run the nvidia-smi command and capture the output + _ = subprocess.check_output(["nvidia-smi"]) + # If the command ran successfully, assume a CUDA device is present + logging.info("A CUDA-capable device is detected.") + except subprocess.CalledProcessError as e: + # If the command fails, it means no CUDA device is detected + logging.info("No CUDA-capable device is detected. Error:", e) + raise Exception("No CUDA-capable device is detected.") command = "/openpose/build/examples/openpose/openpose.bin\ --video {video_path}\ --display 0\