Upload folder using huggingface_hub

heatmap.py

@@ -9,6 +9,9 @@ model_type = "DPT_Large"     # MiDaS v3 - Large     (highest accuracy, slowest i
 
 midas = torch.hub.load("intel-isl/MiDaS", model_type)
 midas_transforms = torch.hub.load("intel-isl/MiDaS", "transforms")
+# Move to CUDA if available
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+midas = midas.to(device)
 
 if model_type == "DPT_Large" or model_type == "DPT_Hybrid":
     transform = midas_transforms.dpt_transform
@@ -18,7 +21,7 @@ else:
 def get_heatmap(src_img):
     cv2image = cv2.imread(src_img)
     img = cv2.cvtColor(cv2image, cv2.COLOR_BGR2RGB)
-    input_batch = transform(img)
+    input_batch = transform(img).to(device)
     with torch.inference_mode():
         prediction = midas(input_batch)