Update app.py

app.py

@@ -40,39 +40,78 @@ def resize_image_pil(image, new_width, new_height):
 
     return resized
 
-def inference(input_img, transparency):
+# def inference(input_img, transparency):
+#     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)
+#     target_layers = [model.layer2[-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 classes[prediction[0].item()], visualization
+
+def inference(input_img, transparency=0.5, target_layer_number=-1):
+    input_img = resize_image_pil(input_img, 32, 32)
+    input_img = np.array(input_img)
+    org_img= input_img
+    
+    input_img = input_img.reshape((32, 32, 3))
+    
     transform = transforms.ToTensor()
     input_img = transform(input_img)
-    input_img = input_img.to(device)
+    
     input_img = input_img.unsqueeze(0)
     outputs = model(input_img)
+
+    softmax = torch.nn.Softmax(dim=0)
+    o = softmax(outputs.flatten())
+    confidences = {classes[i] : float(o[i]) for i in range(10)}
     _, prediction = torch.max(outputs, 1)
-    target_layers = [model.layer2[-2]]
-    cam = GradCAM(model=model, target_layers=target_layers, use_cuda=True)
-    grayscale_cam = cam(input_tensor=input_img, targets=targets)
+    target_layers = [model.layer2[target_layer_number]]
+    cam = GradCAM(model=model, target_layers = target_layers)
+    grayscale_cam = cam(input_tensor=input_img, targets=None)
     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 classes[prediction[0].item()], visualization
+    visualization = show_cam_on_image(
+        org_img/255,
+        grayscale_cam,
+        use_rgb=True,
+        image_weight=transparency
+    )
+
+    return classes[prediction[0].item()], visualization, confidences
+    
+
+
 
 demo = gr.Interface(
-    inference,
-    inputs = [
+    fn=inference,
+    inputs=[
         gr.Image(width=256, height=256, label="Input Image"),
-        gr.Slider(0, 1, value=0.5, label="Overall opacity fo the overlay"),
-        gr.Slider(-2, -1, value=-2, step=1, label="Which GradCAM layer?")    
+        gr.Slider(0,1, value=0.5, label="Overall opacity value"),
+        gr.Slider(-2, -1, value=-2, label="Which model layer to use for GradCAM?")
     ],
     outputs = [
         "text",
         gr.Image(width=256, height=256, label="Output"),
         gr.Label(num_top_classes=3)
     ],
-    title="CIFAR10 trained on ResNet18 with GradCAM feature",
-    description = "A simple Gradio app for checking GradCAM outputs from results of ResNet18 model.",
-    examples = [["cat.jpg", 0.5, -1], ["dog.jpg", 0.7, -2]]
+
+    title="CIFAR10 trained on ResNet18 with GradCAM",
+    
+    description = "A simple Gradio interface to infer on ResNet model with GradCAM results shown on top.",
+    
+    examples = [
+    ["cat.jpg", 0.5, -1],
+    ["dog.jpg", 0.7, -2]
+]
 )
 
 demo.launch()