Deploying to HF

app.py

@@ -32,6 +32,7 @@ class Explainer:
         )
 
         self.transformed_img = transform(img)
+
         self.input = transform_normalize(self.transformed_img)
         self.input = self.input.unsqueeze(0)
 
@@ -46,7 +47,10 @@ class Explainer:
         self.fig_title = 'Predicted: ' + self.pred_label + ' (' + str(round(self.pred_score.squeeze().item(), 2)) + ')'
 
     def convert_fig_to_pil(self, fig):
-        return PIL.Image.frombytes('RGB', fig.canvas.get_width_height(), fig.canvas.tostring_rgb())
+        fig.canvas.draw()
+        data = np.fromstring(fig.canvas.tostring_rgb(), dtype=np.uint8, sep='')
+        data = data.reshape(fig.canvas.get_width_height()[::-1] + (3,))
+        return PIL.Image.fromarray(data)
 
     def shap(self):
         gradient_shap = GradientShap(self.model)
@@ -75,7 +79,7 @@ labels = [ "benign", "malignant", "normal" ]
 def predict(img):
     explainer = Explainer(model, img, labels)
     shap_img = explainer.shap()
-    return [explainer.confidences, shap_img]
+    return [explainer.confidences, shap_img] 
 
 ui = gr.Interface(fn=predict, 
                 inputs=gr.Image(type="pil"),