Update utils.py

utils.py

@@ -9,6 +9,7 @@ from pytorch_grad_cam import GradCAM
 from pytorch_grad_cam.utils.image import show_cam_on_image
 import matplotlib.pyplot as plt
 import textwrap
+import io
 
 def apply_normalization(chennels):
       return nn.BatchNorm2d(chennels)
@@ -88,6 +89,19 @@ def resize_image(image, target_size=(200, 200)):
 def wrap_text(text, width=20):
     return textwrap.fill(text, width)
     
+import io
+# define a function which returns an image as numpy array from figure
+def get_img_from_fig(fig, dpi=180):
+    buf = io.BytesIO()
+    fig.savefig(buf, format="png", dpi=dpi)
+    buf.seek(0)
+    img_arr = np.frombuffer(buf.getvalue(), dtype=np.uint8)
+    buf.close()
+    img = cv2.imdecode(img_arr, 1)
+    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
+
+    return img
+    
 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
@@ -107,9 +121,9 @@ def save_plot_as_image(images,texts, output_path):
         else:
             ax.axis('off')
     plt.tight_layout()
-    # plt.savefig(output_path)
+    # plt.savefig("tmp_arrays.png")
     # plt.close()
-    return plt
+    return get_img_from_fig(plt)
 
 
 # Function to run inference and return top classes
@@ -156,38 +170,6 @@ def get_gradcam(model,input_img, opacity,layer):
     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):
     if show_misclassified:
         return cv2.imread(f"missclassified_images_examples/{int(num)}.png")
@@ -208,27 +190,4 @@ def main_inference(num_of_output_classes,classes,model,input_img):
     confidences = {classes[i]:float(out[i]) for i in range(num_of_output_classes)}
     outputs = model(input_img)
     _, prediction = torch.max(outputs, 1)
-    return confidences
-# def run_inference(input_img, num_of_output_classes,transparency):
-#     transform = transforms.ToTensor()
-#     input_img = transform(input_img)
-#     input_img = input_img.to(device)
-#     input_img = input_img.unsqueeze(0)
-#     softmax = torch.nn.Softmax(dim=0)
-#     outputs = model(input_img)
-#     out = softmax(outputs.flatten())
-#     _, prediction = torch.max(outputs, 1)
-#     confidences = {classes[i]:float(out[i]) for i in range(num_of_output_classes)}
-#     target_layers = [model.convlayer3[-2]]
-    
-#     cam = GradCAM(model=model, target_layers=target_layers, 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=transparency)
-#     return confidences, rgb_img, transparency,grayscale_cam
-
-
+    return confidences