Fix: Inference issue

inference.py

@@ -16,29 +16,24 @@ from pytorch_grad_cam.utils.image import show_cam_on_image
 from pytorch_grad_cam.utils.model_targets import ClassifierOutputTarget
 
 
-def inference(input_img,
-              model,
-              classes,
-              transparency=0.5,
-              number_of_top_classes=3,
-              target_layer_number=4):
+def inference(image, alpha, top_k, target_layer, model=None, classes=None):
     """
     Function to run inference on the input image
-    :param input_img: Image provided by the user
+    :param image: Image provided by the user
+    :param alpha: Percentage of cam overlap over the input image
+    :param top_k: Number of top predictions for the input image
+    :param target_layer: Layer for which GradCam to be shown
     :param model: Model to use for inference
     :param classes: Classes to use for inference
-    :param transparency: Percentage of cam overlap over the input image
-    :param number_of_top_classes: Number of top predictions for the input image
-    :param target_layer_number: Layer for which GradCam to be shown
     """
     # Save a copy of input img
-    org_img = input_img.copy()
+    org_img = image.copy()
 
     # Calculate mean over each channel of input image
-    mean_r, mean_g, mean_b = np.mean(input_img[:, :, 0]/255.), np.mean(input_img[:, :, 1]/255.), np.mean(input_img[:, :, 2]/255.)
+    mean_r, mean_g, mean_b = np.mean(image[:, :, 0]/255.), np.mean(image[:, :, 1]/255.), np.mean(image[:, :, 2]/255.)
 
     # Calculate Standard deviation over each channel
-    std_r, std_g, std_b = np.std(input_img[:, :, 0]/255.), np.std(input_img[:, :, 1]/255.), np.std(input_img[:, :, 2]/255.)
+    std_r, std_g, std_b = np.std(image[:, :, 0]/255.), np.std(image[:, :, 1]/255.), np.std(image[:, :, 2]/255.)
 
     # Convert img to tensor and normalize it
     _transform = transforms.Compose([
@@ -47,7 +42,7 @@ def inference(input_img,
         ])
 
     # Preprocess the input image
-    input_tensor = _transform(input_img)
+    input_tensor = _transform(image)
 
     # Create a mini-batch as expected by the model
     input_tensor = input_tensor.unsqueeze(0)
@@ -66,7 +61,7 @@ def inference(input_img,
 
     # Select the top classes based on user input
     sorted_confidences = sorted(confidences.items(), key=lambda val: val[1], reverse=True)
-    show_confidences = OrderedDict(sorted_confidences[:number_of_top_classes])
+    show_confidences = OrderedDict(sorted_confidences[:top_k])
 
     # Map layer numbers to meaningful parts of the ResNet architecture
     _layers = {
@@ -79,8 +74,8 @@ def inference(input_img,
     }
 
     # Ensure valid layer selection
-    target_layer_number = min(max(target_layer_number, 1), 6)
-    target_layers = [_layers[target_layer_number]]
+    target_layer = min(max(target_layer, 1), 6)
+    target_layers = [_layers[target_layer]]
 
     # Get the class activations from the selected layer
     cam = GradCAM(model=model, target_layers=target_layers)
@@ -98,5 +93,5 @@ def inference(input_img,
     grayscale_cam = grayscale_cam[0, :]
 
     # Overlay input image with Class activations
-    visualization = show_cam_on_image(org_img/255., grayscale_cam, use_rgb=True, image_weight=transparency)
+    visualization = show_cam_on_image(org_img/255., grayscale_cam, use_rgb=True, image_weight=alpha)
     return show_confidences, visualization