added support for running demos on CPU when no CUDA devices are found

.gitignore

@@ -0,0 +1,23 @@
+
+images/
+
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# PyTorch weights
+*.tar
+*.pth
+*.gz
+Untitled.ipynb
+Testing notebook.ipynb

image_demo.py

@@ -18,7 +18,8 @@
 
 def main():
     model = posenet.load_model(args.model)
-    model = model.cuda()
+    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    model = model.to(device=device)
     output_stride = model.output_stride
 
     if args.output_dir:
@@ -34,7 +35,7 @@ def main():
             f, scale_factor=args.scale_factor, output_stride=output_stride)
 
         with torch.no_grad():
-            input_image = torch.Tensor(input_image).cuda()
+            input_image = torch.Tensor(input_image).to(device=device)
 
             heatmaps_result, offsets_result, displacement_fwd_result, displacement_bwd_result = model(input_image)
 

webcam_demo.py

@@ -16,7 +16,8 @@
 
 def main():
     model = posenet.load_model(args.model)
-    model = model.cuda()
+    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    model = model.to(device=device)
     output_stride = model.output_stride
 
     cap = cv2.VideoCapture(args.cam_id)
@@ -30,7 +31,7 @@ def main():
             cap, scale_factor=args.scale_factor, output_stride=output_stride)
 
         with torch.no_grad():
-            input_image = torch.Tensor(input_image).cuda()
+            input_image = torch.Tensor(input_image).to(device=device)
 
             heatmaps_result, offsets_result, displacement_fwd_result, displacement_bwd_result = model(input_image)
 
@@ -59,4 +60,4 @@ def main():
 
 
 if __name__ == "__main__":
-    main()
+    main()