Update utils.py

utils.py

@@ -7,6 +7,8 @@ import numpy as np
 from pytorch_grad_cam import GradCAM
 from pytorch_grad_cam import GradCAM
 from pytorch_grad_cam.utils.image import show_cam_on_image
+import matplotlib.pyplot as plt
+import textwrap
 
 def apply_normalization(chennels):
       return nn.BatchNorm2d(chennels)
@@ -78,7 +80,38 @@ class CustomResnet(nn.Module):
         x = x.view(-1, 512)
         x = self.linear1(x)
         return F.log_softmax(x, dim=-1)
+
+
+def resize_image(image, target_size=(200, 200)):
+    return cv2.resize(image, target_size)
+
+def wrap_text(text, width=20):
+    return textwrap.fill(text, width)
     
+def save_plot_as_image(images,texts, output_path):
+    num_images = len(images)
+    num_cols = min(4, num_images)  # Assuming you want a maximum of 4 columns
+    num_rows = (num_images - 1) // num_cols + 1
+
+    fig, axes = plt.subplots(num_rows, num_cols, figsize=(3 * num_cols, 3 * num_rows))
+
+    subplot_height = 0.9 / num_rows  # Adjust this value to control the height of each subplot
+    plt.subplots_adjust(hspace=subplot_height)
+    for i, ax in enumerate(axes.flat):
+        if i < num_images:
+            ax.imshow(images[i], cmap='gray')
+            ax.axis('off')
+            if texts is not None and i < len(texts):
+                wrapped_text = wrap_text(texts[i])
+                ax.set_title(wrapped_text, fontsize=12, pad=5)
+        else:
+            ax.axis('off')
+    plt.tight_layout()
+    # plt.savefig(output_path)
+    # plt.close()
+    return plt
+
+
 # Function to run inference and return top classes
 def get_gradcam(model,input_img, opacity,layer):
     targets = None
@@ -92,23 +125,67 @@ def get_gradcam(model,input_img, opacity,layer):
     input_img = input_img.to(device)
     input_img = input_img.unsqueeze(0)
     outputs = model(input_img)
-    _, prediction = torch.max(outputs, 1)
-    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, :]
-    img = input_img.squeeze(0).to('cpu')
-    img = inv_normalize(img)
-    rgb_img = np.transpose(img, (1, 2, 0))
-    rgb_img = rgb_img.numpy()
-    visualization = show_cam_on_image(rgb_img, grayscale_cam, use_rgb=True, image_weight=opacity)
-    return visualization
+    if layer == "convblock1":
+        target_layers = model.convlayer1
+    elif layer == "convblock2":
+        target_layers = model.convlayer2
+    elif layer == "resblock1":
+        target_layers = model.reslayer1
+    elif layer == "resblock2":
+        target_layers = model.reslayer2
+    elif layer == "convblock3":
+        target_layers = model.convlayer3
+    
+    layer_to_user = []
+    for i in target_layers:
+        if str(i) != "ReLU()":
+            layer_to_user.append(i)
+    print(layer_to_user)
+    final_outputs,texts = [],[]
+    for i in range(len(layer_to_user)):
+        cam = GradCAM(model=model, target_layers=[layer_to_user[i]], use_cuda=True)
+        grayscale_cam = cam(input_tensor=input_img, targets=targets)
+        grayscale_cam = grayscale_cam[0, :]
+        img = input_img.squeeze(0).to('cpu')
+        img = inv_normalize(img)
+        rgb_img = np.transpose(img, (1, 2, 0))
+        rgb_img = rgb_img.numpy()
+        visualization = show_cam_on_image(rgb_img, grayscale_cam, use_rgb=True, image_weight=opacity)
+        final_outputs.append(resize_image(visualization))
+        texts.append(str(layer_to_user[i]))
+    figure = save_plot_as_image(final_outputs,texts, "plot.png")
+    return figure
+    
+# # Function to run inference and return top classes
+# 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],
+#         std=[1/0.23, 1/0.23, 1/0.23]
+#     )
+#     device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+#     transform = transforms.ToTensor()
+#     input_img = transform(input_img)
+#     input_img = input_img.to(device)
+#     input_img = input_img.unsqueeze(0)
+#     outputs = model(input_img)
+#     _, prediction = torch.max(outputs, 1)
+#     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, :]
+#     img = input_img.squeeze(0).to('cpu')
+#     img = inv_normalize(img)
+#     rgb_img = np.transpose(img, (1, 2, 0))
+#     rgb_img = rgb_img.numpy()
+#     visualization = show_cam_on_image(rgb_img, grayscale_cam, use_rgb=True, image_weight=opacity)
+#     return visualization
 
 
 def get_misclassified_images(show_misclassified,num):