Update app.py

app.py

@@ -3,7 +3,8 @@ import cv2
 from utils.detect import create_mtcnn_net, MtcnnDetector
 from utils.vision import vis_face
 import argparse
-
+from PIL import Image
+import numpy as np
 
 MIN_FACE_SIZE = 3
 
@@ -33,7 +34,18 @@ def greet(name):
     img_bg = cv2.merge([r,g,b])
     p_bboxs, r_bboxs, bboxs, landmarks = mtcnn_detector.detect_face(img)
     save_name = args.save_name
-    return vis_face(img_bg, bboxs, landmarks, MIN_FACE_SIZE, save_name)
+    fig = vis_face(img_bg, bboxs, landmarks, MIN_FACE_SIZE, save_name)
+    fig.canvas.draw()
+
+    # Get the RGB buffer from the figure
+    w, h = fig.canvas.get_width_height()
+    buf = np.fromstring(fig.canvas.tostring_rgb(), dtype=np.uint8)
+    buf.shape = (h, w, 3)
+
+    # canvas.tostring_rgb give pixmap in RGB mode. 
+    # Roll the ALPHA channel to have it in RGBA mode
+    buf = np.roll(buf, 3, axis=2)
+    return Image.fromarray(buf)
 iface = gr.Interface(fn=greet, 
                     inputs=gr.Image(type="filepath"),
                     outputs="image")