up

app.py

@@ -3,6 +3,15 @@ import numpy as np
 from models import make_inpainting
 import utils
 
+from transformers import MaskFormerImageProcessor, MaskFormerForInstanceSegmentation
+from PIL import Image
+import requests
+from transformers import pipeline
+import torch
+import random
+import io
+import base64
+
 def removeFurniture(input_img1,
             input_img2,
             positive_prompt,
@@ -36,8 +45,24 @@ def removeFurniture(input_img1,
 
     return retList
 
-def segmentation(image):
-    return image
+def imageToString(img):
+
+    output = io.BytesIO()
+    img.save(output, format="png")
+    return output.getvalue()
+
+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)
+    for p in results:
+        p['mask'] = utils.image_to_byte_array(p['mask'])
+        p['mask'] = base64.b64encode(p['mask'])
+    #print(results)
+    return str(results)
+    
 
 def upscale(image):
     return image
@@ -66,7 +91,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.Image())
+            gr.Button("Segmentation").click(segmentation, inputs=gr.Image(type="pil"), outputs=gr.TextArea())
         with gr.Column():
             gr.Button("Upscale").click(upscale, inputs=gr.Image(type="pil"), outputs=gr.Image())
 

pipelines.py

@@ -2,8 +2,6 @@ import logging
 import torch
 import time
 from diffusers import StableDiffusionInpaintPipeline
-
-# from config import WIDTH, HEIGHT
 from helpers import flush
 
 LOGGING = logging.getLogger(__name__)