Changed from plot output to img ?

app.py

@@ -139,7 +139,12 @@ def inference(query, labels, n_supp=10,
 
         probs = output.softmax(dim=-1).detach().cpu().numpy()
 
-        return {k: float(v) for k, v in zip(labels, probs[0, 0])}, fig
+        # it used to be a matplotlib figure, turning it into a PIL image
+        fig.canvas.draw()
+        w, h = fig.canvas.get_width_height()
+        img = np.frombuffer(fig.canvas.tostring_rgb(), dtype=np.uint8).reshape(h, w, 3)
+        return {k: float(v) for k, v in zip(labels, probs[0, 0])}, Image.fromarray(img)
+        #return {k: float(v) for k, v in zip(labels, probs[0, 0])}, fig
 
 
 # DEBUG