up

app.py

@@ -11,6 +11,7 @@ import torch
 import random
 import io
 import base64
+import json
 
 def removeFurniture(input_img1,
             input_img2,
@@ -55,13 +56,13 @@ def segmentation(img):
     print("segmentation")
 
     # semantic_segmentation = pipeline("image-segmentation", "nvidia/segformer-b1-finetuned-cityscapes-1024-1024")
-    semantic_segmentation = pipeline("image-segmentation", "facebook/maskformer-swin-large-ade")    
-    results = semantic_segmentation(img)
+    pipe = pipeline("image-segmentation", "facebook/maskformer-swin-large-ade")    
+    results = pipe(img)
     for p in results:
         p['mask'] = utils.image_to_byte_array(p['mask'])
-        p['mask'] = base64.b64encode(p['mask'])
+        p['mask'] = base64.b64encode(p['mask']).decode("utf-8")
     #print(results)
-    return str(results)
+    return json.dumps(results)
     
 
 def upscale(image):
@@ -91,7 +92,7 @@ with gr.Blocks() as app:
                                                 gr.Image(),
                                                 gr.Image()])
         with gr.Column():  
-            gr.Button("Segmentation").click(segmentation, inputs=gr.Image(type="pil"), outputs=gr.TextArea())
+            gr.Button("Segmentation").click(segmentation, inputs=gr.Image(type="pil"), outputs=gr.JSON())
         with gr.Column():
             gr.Button("Upscale").click(upscale, inputs=gr.Image(type="pil"), outputs=gr.Image())