Update utils.py

utils.py

@@ -80,7 +80,7 @@ class CustomResnet(nn.Module):
         return F.log_softmax(x, dim=-1)
     
 # Function to run inference and return top classes
-def get_gradcam(model,input_img, opacity):
+def get_gradcam(model,input_img, opacity,layer):
     targets = None
     inv_normalize = transforms.Normalize(
         mean=[-0.50/0.23, -0.50/0.23, -0.50/0.23],
@@ -93,7 +93,13 @@ def get_gradcam(model,input_img, opacity):
     input_img = input_img.unsqueeze(0)
     outputs = model(input_img)
     _, prediction = torch.max(outputs, 1)
-    target_layers = [model.convlayer3[-2]]
+    if layer == "layer3":
+        target_layers = [model.convlayer3[-2]]
+    if layer == "layer2":
+        target_layers = [model.convlayer2[-2]]
+    if layer == "layer1":
+        target_layers = [model.convlayer1[-2]]
+    #target_layers = [model.convlayer3[-2]]
     cam = GradCAM(model=model, target_layers=target_layers, use_cuda=False)
     grayscale_cam = cam(input_tensor=input_img, targets=targets)
     grayscale_cam = grayscale_cam[0, :]