Upload folder using huggingface_hub

demo.py

@@ -1,5 +1,5 @@
 import os
-os.system("pip install datasets einops cupy-cuda11x tabulate opencv-python -U")
+os.system("pip install datasets einops tabulate opencv-python ttach -U")
 
 import os, cv2, time, math
 print("=> Loading libraries...")
@@ -15,6 +15,8 @@ from pytorch_grad_cam import GradCAM
 from pytorch_grad_cam.utils.image import show_cam_on_image
 
 
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
 print(f"=> Libraries loaded in {time.time()- start:.2f} sec(s).")
 print("=> Loading model...")
 start = time.time()
@@ -25,7 +27,7 @@ crop_pct = 0.9
 IMAGENET_DEFAULT_MEAN = (0.485, 0.456, 0.406)
 IMAGENET_DEFAULT_STD = (0.229, 0.224, 0.225)
 
-model = create_model(f"tpmlp_{size}").cuda()
+model = create_model(f"tpmlp_{size}").to(device)
 load_checkpoint(model, f"../tpmlp_{size}.pth.tar", True)
 model.eval()
 
@@ -56,7 +58,7 @@ def transform(img):
 def predict(inp):
     img, inp = transform(inp)
     inp = inp.unsqueeze(0)
-    with GradCAM(model=model, target_layers=[model.layers[3]], use_cuda=True) as cam:
+    with GradCAM(model=model, target_layers=[model.layers[3]], use_cuda=device=="cuda") as cam:
         grayscale_cam, probs = cam(input_tensor=inp, aug_smooth=False, eigen_smooth=False, return_probs=True)
         
         # Here grayscale_cam has only one image in the batch