Deploying to HF

app.py

@@ -1,5 +1,5 @@
 import PIL
-from captum.attr import GradientShap
+from captum.attr import GradientShap, Occlusion, LayerGradCam, LayerAttribution, IntegratedGradients
 from captum.attr import visualization as viz
 import torch
 from torchvision import transforms
@@ -65,6 +65,58 @@ class Explainer:
         fig.suptitle(self.fig_title, fontsize=12)
         return self.convert_fig_to_pil(fig)
 
+    def occlusion(self):
+        occlusion = Occlusion(model)
+
+        attributions_occ = occlusion.attribute(self.input,
+                                               target=self.pred_label_idx,
+                                               strides=(3, 8, 8),
+                                               sliding_window_shapes=(3,15, 15),
+                                               baselines=0)
+
+        fig, _ = viz.visualize_image_attr_multiple(np.transpose(attributions_occ.squeeze().cpu().detach().numpy(), (1,2,0)),
+                                            np.transpose(self.transformed_img.squeeze().cpu().detach().numpy(), (1,2,0)),
+                                            ["original_image", "heat_map", "heat_map", "masked_image"],
+                                            ["all", "positive", "negative", "positive"],
+                                            show_colorbar=True,
+                                            titles=["Original", "Positive Attribution", "Negative Attribution", "Masked"],
+                                            fig_size=(18, 6)
+                                            )
+        fig.suptitle(self.fig_title, fontsize=12)
+        return self.convert_fig_to_pil(fig)
+    
+    def gradcam(self):
+        layer_gradcam = LayerGradCam(self.model, self.model.layer3[1].conv2)
+        attributions_lgc = layer_gradcam.attribute(self.input, target=self.pred_label_idx)
+
+        #_ = viz.visualize_image_attr(attributions_lgc[0].cpu().permute(1,2,0).detach().numpy(),
+        #                            sign="all",
+        #                            title="Layer 3 Block 1 Conv 2")
+        upsamp_attr_lgc = LayerAttribution.interpolate(attributions_lgc, self.input.shape[2:])
+
+        fig, _ = viz.visualize_image_attr_multiple(upsamp_attr_lgc[0].cpu().permute(1,2,0).detach().numpy(),
+                                            self.transformed_img.permute(1,2,0).numpy(),
+                                            ["original_image","blended_heat_map","masked_image"],
+                                            ["all","positive","positive"],
+                                            show_colorbar=True,
+                                            titles=["Original", "Positive Attribution", "Masked"],
+                                            fig_size=(18, 6))
+        return self.convert_fig_to_pil(fig)
+    
+    def integrated_gradients(self):
+        integrated_gradients = IntegratedGradients(self.model)
+        attributions_ig = integrated_gradients.attribute(self.input, target=self.pred_label_idx, n_steps=50)
+
+        fig, _ = viz.visualize_image_attr_multiple(np.transpose(attributions_ig.squeeze().cpu().detach().numpy(), (1,2,0)),
+                                            np.transpose(self.transformed_img.squeeze().cpu().detach().numpy(), (1,2,0)),
+                                            ["original_image", "heat_map", "masked_image"],
+                                            ["all", "positive", "positive"],
+                                            show_colorbar=True,
+                                            titles=["Original", "Attribution", "Masked"],
+                                            fig_size=(18, 6))
+        fig.suptitle(self.fig_title, fontsize=12)
+        return self.convert_fig_to_pil(fig)
+
 def create_model_from_checkpoint():
     # Loads a model from a checkpoint
     model = resnet50()
@@ -78,11 +130,10 @@ labels = [ "benign", "malignant", "normal" ]
 
 def predict(img):
     explainer = Explainer(model, img, labels)
-    shap_img = explainer.shap()
-    return [explainer.confidences, shap_img] 
+    return [explainer.confidences, explainer.shap(), explainer.occlusion(), explainer.gradcam(), explainer.integrated_gradients()] 
 
 ui = gr.Interface(fn=predict, 
                 inputs=gr.Image(type="pil"),
-                outputs=[gr.Label(num_top_classes=3), gr.Image(type="pil")],
+                outputs=[gr.Label(num_top_classes=3), gr.Image(type="pil"), gr.Image(type="pil"), gr.Image(type="pil"), gr.Image(type="pil")],
                 examples=["benign (52).png", "benign (243).png", "malignant (127).png", "malignant (201).png", "normal (81).png", "normal (101).png"]).launch()
 ui.launch(share=True)