Update app.py

app.py

@@ -48,17 +48,24 @@ def plot_outlines(img, masks):
         for o in outpix:
             ax.plot(o[:,0], img.shape[0]-o[:,1], color=[1,0,0], lw=1)
     ax.axis('off')
-    bytes_image = io.BytesIO()
-    plt.savefig(bytes_image, format='png', facecolor=fig.get_facecolor(), edgecolor='none')
-    bytes_image.seek(0)
-    img_arr = np.frombuffer(bytes_image.getvalue(), dtype=np.uint8)
-    bytes_image.close()
-    img = cv2.imdecode(img_arr, 1)
-    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
-    del bytes_image
-    fig.clf()
-    plt.close(fig)
-    return img
+    
+    #bytes_image = io.BytesIO()
+    #plt.savefig(bytes_image, format='png', facecolor=fig.get_facecolor(), edgecolor='none')
+    #bytes_image.seek(0)
+    #img_arr = np.frombuffer(bytes_image.getvalue(), dtype=np.uint8)
+    #bytes_image.close()
+    #img = cv2.imdecode(img_arr, 1)
+    #img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
+    #del bytes_image
+    #fig.clf()
+    #plt.close(fig)
+
+    img = plt.gcf()
+    buf = io.BytesIO()
+    figure.savefig(buf, bbox_inches='tight')
+    buf.seek(0)
+    output_pil_img = Image.open(buf)
+    return output_pil_img
 
 def plot_overlay(img, masks):
     img = normalize99(img.astype(np.float32).mean(axis=-1))
@@ -137,7 +144,7 @@ with gr.Blocks(title = "Hello",
 
         with gr.Column(scale=1):            
             img_overlay = gr.Image(label = "Output image", type = "numpy")
-            img_outlines = gr.Image(label = "Output image", type = "numpy")            
+            img_outlines = gr.Image(label = "Output image", type = "pil")            
             flows = gr.Image(label = "Output image", type = "numpy")
             masks = gr.Image(label = "Output image", type = "numpy")