Deploying to HF

app.py

@@ -52,10 +52,10 @@ class Explainer:
         data = data.reshape(fig.canvas.get_width_height()[::-1] + (3,))
         return PIL.Image.fromarray(data)
 
-    def shap(self):
+    def shap(self, n_samples, stdevs):
         gradient_shap = GradientShap(self.model)
         rand_img_dist = torch.cat([self.input * 0, self.input * 1])
-        attributions_gs = gradient_shap.attribute(self.input, n_samples=50, stdevs=0.0001, baselines=rand_img_dist, target=self.pred_label_idx)
+        attributions_gs = gradient_shap.attribute(self.input, n_samples=int(n_samples), stdevs=stdevs, baselines=rand_img_dist, target=self.pred_label_idx)
         fig, _ = viz.visualize_image_attr_multiple(np.transpose(attributions_gs.squeeze().cpu().detach().numpy(), (1,2,0)),
                                             np.transpose(self.transformed_img.squeeze().cpu().detach().numpy(), (1,2,0)),
                                             ["original_image", "heat_map"],
@@ -65,13 +65,13 @@ class Explainer:
         fig.suptitle("SHAP | " + self.fig_title, fontsize=12)
         return self.convert_fig_to_pil(fig)
 
-    def occlusion(self):
+    def occlusion(self, stride, sliding_window):
         occlusion = Occlusion(model)
 
         attributions_occ = occlusion.attribute(self.input,
                                                target=self.pred_label_idx,
-                                               strides=(3, 8, 8),
-                                               sliding_window_shapes=(3,15, 15),
+                                               strides=(3, int(stride), int(stride)),
+                                               sliding_window_shapes=(3, int(sliding_window), int(sliding_window)),
                                                baselines=0)
 
         fig, _ = viz.visualize_image_attr_multiple(np.transpose(attributions_occ.squeeze().cpu().detach().numpy(), (1,2,0)),
@@ -103,20 +103,6 @@ class Explainer:
                                             fig_size=(18, 6))
         fig.suptitle("GradCAM layer3[1].conv2 | " + self.fig_title, fontsize=12)
         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("Integrated gradients | " + self.fig_title, fontsize=12)
-        return self.convert_fig_to_pil(fig)
 
 def create_model_from_checkpoint():
     # Loads a model from a checkpoint
@@ -129,12 +115,21 @@ def create_model_from_checkpoint():
 model = create_model_from_checkpoint()
 labels = [ "benign", "malignant", "normal" ]
 
-def predict(img):
+def predict(img, shap_samples, shap_stdevs, occlusion_stride, occlusion_window):
     explainer = Explainer(model, img, labels)
-    return [explainer.confidences, explainer.shap(), explainer.occlusion(), explainer.gradcam()] 
+    return [explainer.confidences,
+            explainer.shap(shap_samples, shap_stdevs),
+            explainer.occlusion(occlusion_stride, occlusion_window),
+            explainer.gradcam()] 
 
 ui = gr.Interface(fn=predict, 
-                inputs=gr.Image(type="pil"),
+                inputs=[
+                    gr.Image(type="pil"),
+                    gr.Slider(minimum=10, maximum=100, default=50, label="SHAP Samples"),
+                    gr.Slider(minimum=0.0001, maximum=0.01, default=0.0001, label="SHAP Stdevs"),
+                    gr.Slider(minimum=4, maximum=80, default=8, label="Occlusion Stride"),
+                    gr.Slider(minimum=4, maximum=80, default=15, label="Occlusion Window")  
+                ],
                 outputs=[gr.Label(num_top_classes=3), 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)