adding outlines

app.py

@@ -1,10 +1,10 @@
-import os
 import numpy as np
 import gradio as gr
 import cv2
 from cellpose import models
 from matplotlib.colors import hsv_to_rgb
-import os
+import matplotlib.pyplot as plt
+import os, io, base64
 
 try:
     model = models.CellposeModel(gpu=False, pretrained_model="cyto3")
@@ -22,7 +22,7 @@ def plot_flows(y):
     flow = (hsv_to_rgb(HSV) * 255).astype(np.uint8)
     return flow
 
-def plot_outlines(masks):
+def plot_outlines(img, masks):
     outpix = []
     contours, hierarchy = cv2.findContours(masks.astype(np.int32), mode=cv2.RETR_FLOODFILL, method=cv2.CHAIN_APPROX_SIMPLE)
     for c in range(len(contours)):
@@ -31,7 +31,29 @@ def plot_outlines(masks):
             peri = cv2.arcLength(contours[c], True)
             approx = cv2.approxPolyDP(contours[c], 0.001, True)[:,0,:]
             outpix.append(approx)
-    return outpix
+    
+    figsize = (6,6)
+    if img.shape[0]>img.shape[1]:
+        figsize = (6*img.shape[1]/img.shape[0], 6)
+    else:
+        figsize = (6, 6*img.shape[0]/img.shape[1])
+    fig = plt.figure(figsize=figsize, facecolor='k')
+    ax = fig.add_axes([0.0,0.0,1,1])
+    ax.set_xlim([0,img.shape[1]])
+    ax.set_ylim([0,img.shape[0]])
+    ax.imshow(img[::-1], origin='upper')
+    if outpix is not None:
+        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_html = base64.b64encode(bytes_image.getvalue()).decode()
+    del bytes_image
+    fig.clf()
+    plt.close(fig)
+    return img_html
 
 def plot_overlay(img, masks):
     img = normalize99(img.astype(np.float32).mean(axis=-1))
@@ -66,7 +88,7 @@ def image_resize(img, resize=224):
     return img
 
 def cellpose_segment(img):
-    img_input = image_resize(img)
+    img_input = image_resize(img.copy())
     masks, flows, _ = model.eval(img_input, channels=[0,0])
     flows = flows[0]
     # masks = np.zeros(img.shape[:2])
@@ -77,16 +99,16 @@ def cellpose_segment(img):
         masks = cv2.resize(masks.astype('uint16'), target_size, interpolation=cv2.INTER_NEAREST).astype('uint16')
         flows = cv2.resize(flows.astype('float32'), target_size).astype('uint8')
 
-    outpix = plot_outlines(masks)
+    outpix = plot_outlines(img, masks)
     overlay = plot_overlay(img, masks)
 
-    return overlay, flows, masks 
+    return outpix, overlay, flows, masks 
 
 # Gradio Interface
 iface = gr.Interface(
     fn=cellpose_segment, 
     inputs="image", 
-    outputs=["image", "image", "image"],
+    outputs=["image", "image", "image", "image"],
     title="cellpose segmentation",
     description="upload an image, then cellpose will segment it at a max size of 224x224"
 )