Update app.py

app.py

@@ -40,6 +40,40 @@ def inference(input_img, transparency = 0.5, target_layer_number = -1):
     visualization = show_cam_on_image(org_img/255, grayscale_cam, use_rgb=True, image_weight=transparency)
     return confidences, visualization
 
+def inference_confidences(input_img, transparency = 0.5, target_layer_number = -1):
+    transform = transforms.ToTensor()
+    org_img = input_img
+    input_img = transform(input_img)
+    input_img = input_img
+    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)}
+    return confidences
+
+def inference_visualization(input_img, transparency = 0.5, target_layer_number = -1):
+    transform = transforms.ToTensor()
+    org_img = input_img
+    input_img = transform(input_img)
+    input_img = input_img
+    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[target_layer_number]]
+    cam = GradCAM(model=model, target_layers=target_layers, use_cuda=False)
+    grayscale_cam = cam(input_tensor=input_img, targets=None)
+    grayscale_cam = grayscale_cam[0, :]
+    img = input_img.squeeze(0)
+    img = inv_normalize(img)
+    rgb_img = np.transpose(img, (1, 2, 0))
+    rgb_img = rgb_img.numpy()
+    visualization = show_cam_on_image(org_img/255, grayscale_cam, use_rgb=True, image_weight=transparency)
+    return visualization
+
 title = "CIFAR10 trained on ResNet18 Model with GradCAM"
 description = "Gradio interface to infer on ResNet18 model, and get GradCAM results"
 examples = [["cat.jpg", 0.5, -1], ["dog.jpg", 0.5, -1]]
@@ -62,7 +96,20 @@ demo = gr.Interface(
         gr.Image(label="Uploaded image" if "Interface" in gr.__file__ else "Image"),
         gr.Button(label="Submit", type="boolean")
     ],
-    outputs=gr.outputs.Textbox(label="Results"),
+    outputs = [gr.Label(num_top_classes=3), gr.Image(shape=(32, 32), label="Output").style(width=128, height=128)],
+    examples = examples,
     live=True
 )
-demo.launch()
+
+# Callback function for the Gradio interface
+def gradio_callback(view_gradcam, num_gradcam_images, layer_name, opacity,
+                    view_misclassified, num_misclassified_images,
+                    image_source, uploaded_image, submit,input_img, transparency = 0.5, target_layer_number = -1):
+    confidence = inference_confidences(input_img, transparency = 0.5, target_layer_number = -1)
+    visualization = inference_visualization(input_img, transparency = 0.5, target_layer_number = -1)
+    return confidence, visualization
+
+
+# Set the callback function to the Gradio interface
+demo.fn = gradio_callback
+demo.launch()