enable dynamic selection of CUDA GPU only if available, else GPU inference

lungtumormask/__main__.py

@@ -17,4 +17,4 @@ def main():
     argsin = sys.argv[1:]
     args = parser.parse_args(argsin)
 
-    mask.mask(args.input, args.output)
+    mask.mask(args.input, args.output)

lungtumormask/dataprocessing.py

@@ -10,7 +10,8 @@ from monai.transforms import (Compose, LoadImaged, ToNumpyd, ThresholdIntensityd
 
 def mask_lung(scan_path, batch_size=20):
     model = lungmask.mask.get_model('unet', 'R231')
-    device = torch.device('cuda')
+    if T.cuda.is_available():
+        device = torch.device('cuda')
     model.to(device)
 
     scan_dict = {

lungtumormask/mask.py

@@ -5,10 +5,14 @@ import torch as T
 import nibabel
 
 def load_model():
+    if T.cuda.is_available():
+        gpu_device = T.device('cuda')
+    else:
+        gpu_device = T.device('cpu')
     model = UNet_double(3, 1, 1, tuple([64, 128, 256, 512, 1024]), tuple([2 for i in range(4)]), num_res_units = 0)
-    state_dict = T.hub.load_state_dict_from_url("https://github.com/VemundFredriksen/LungTumorMask/releases/download/0.0/dc_student.pth", progress=True, map_location=T.device('cuda:0'))
+    state_dict = T.hub.load_state_dict_from_url("https://github.com/VemundFredriksen/LungTumorMask/releases/download/0.0/dc_student.pth", progress=True, map_location=gpu_device)
     #model.load_state_dict(T.load("D:\\OneDrive\\Skole\\Universitet\\10. Semester\\Masteroppgave\\bruk_for_full_model.pth", map_location="cuda:0"))
-    model.load_state_dict(state_dict) 
+    model.load_state_dict(state_dict)
     model.eval()
     return model