Changed Code for CUDA

Model_Class.py

@@ -59,14 +59,15 @@ val_transforms_416x628 = Compose(
     ]
 )
 
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 checkpoint = torch.load("classification_model.ckpt", map_location=torch.device('cpu'))
-model = ResNet().to(device)
+model = ResNet()
 model.load_state_dict(checkpoint["state_dict"])
 model.eval()
 
 
-def load_and_classify_image(image_path):
+def load_and_classify_image(image_path, device):
+
+    model = model.to(device)
     image = val_transforms_416x628(image_path)
     image = image.unsqueeze(0).to(device)
 
@@ -76,8 +77,10 @@ def load_and_classify_image(image_path):
         return prediction.to('cpu'), image.to('cpu')
     
 
-def make_GradCAM(image):
+def make_GradCAM(image, device):
 
+    model = model.to(device)
+    image = image.to(device)
     model.eval()
     target_layers = [model.model.layer4[-1]]
 
@@ -90,7 +93,7 @@ def make_GradCAM(image):
         aug_smooth=False,
         eigen_smooth=True,
     )
-    grayscale_cam = grayscale_cam.squeeze()
+    grayscale_cam = grayscale_cam.to('cpu').squeeze()
 
     jet = plt.colormaps.get_cmap("inferno")
     newcolors = jet(np.linspace(0, 1, 256))

Model_Seg.py

@@ -39,10 +39,8 @@ model = UNet(
     num_res_units=3
 )
 
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-
 checkpoint_path = 'segmentation_model.pt'
-checkpoint = torch.load(checkpoint_path, map_location=device)  
+checkpoint = torch.load(checkpoint_path, map_location='cpu')  
 assert model.state_dict().keys() == checkpoint['network'].keys(), "Model and checkpoint keys do not match"
 
 model.load_state_dict(checkpoint['network'])
@@ -73,9 +71,9 @@ post_transforms = Compose([
 
 
 
-def load_and_segment_image(input_image_path):
+def load_and_segment_image(input_image_path, device):
         
-    
+    model = model.to(device)
     image_tensor = pre_transforms(input_image_path)
     image_tensor = image_tensor.unsqueeze(0).to(device)
 

app.py

@@ -7,6 +7,9 @@ import SimpleITK as sitk
 import torch
 from numpy import uint8
 import spaces
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
 image_base64 = utils.image_to_base64("anatomy_aware_pipeline.png")
 article_html = f"<img src='data:image/png;base64,{image_base64}' alt='Anatomical pipeline illustration' style='width:100%;'>"
 
@@ -64,7 +67,7 @@ def predict_image(input_image, input_file):
     else:
         return None , None , "Please input an image before pressing run" , None , None
 
-    image_mask = Model_Seg.load_and_segment_image(image_path)
+    image_mask = Model_Seg.load_and_segment_image(image_path, device)
 
     overlay_image_np, original_image_np = utils.overlay_mask(image_path, image_mask)
 
@@ -75,10 +78,10 @@ def predict_image(input_image, input_file):
     cropped_boxed_array = sitk.GetArrayFromImage(cropped_boxed_im)
     cropped_boxed_array_disp = cropped_boxed_array.squeeze()
     cropped_boxed_tensor = torch.Tensor(cropped_boxed_array)
-    prediction, image_transformed = Model_Class.load_and_classify_image(cropped_boxed_tensor)
+    prediction, image_transformed = Model_Class.load_and_classify_image(cropped_boxed_tensor, device)
 
 
-    gradcam = Model_Class.make_GradCAM(image_transformed)
+    gradcam = Model_Class.make_GradCAM(image_transformed, device)
     
     nr_axSpA_prob = float(prediction[0].item())
     r_axSpA_prob = float(prediction[1].item())