add SAM

app.py

@@ -9,6 +9,30 @@ from PIL import Image
 import numpy as np
 import gradio as gr
 
+
+def inference_mask1_sam(prompt,
+              img,
+              img_):
+
+    files = {
+        "useSam" : 1,
+        "pimage" : resizeImg(prompt["image"]),
+        "pmask" : resizeImg(prompt["mask"]),
+        "img" : resizeImg(img),
+        "img_" : resizeImg(img_)
+    }
+    r = requests.post("http://120.92.79.209/painter/run", json = files)
+    a = json.loads(r.text)
+
+    res = []
+
+    for i in range(len(a)):
+        #out = Image.open(io.BytesIO(base64.b64decode(a[i])))
+        #out = out.resize((224, 224))
+        #res.append(np.uint8(np.array(out)))
+        res.append(np.uint8(np.array(Image.open(io.BytesIO(base64.b64decode(a[i]))))))
+    return res[1:] # remove prompt image
+
 def inference_mask1(prompt,
               img,
               img_):
@@ -56,6 +80,15 @@ examples = [
             ['./images/ydt_2.jpg', './images/ydt_1.jpg', './images/ydt_3.jpg'],
            ]
 
+examples_sam = [
+            ['./images/hmbb_1.jpg', './images/hmbb_2.jpg', './images/hmbb_3.jpg'],
+            ['./images/street_1.jpg', './images/street_2.jpg', './images/street_3.jpg'],
+            ['./images/tom_1.jpg', './images/tom_2.jpg', './images/tom_3.jpg'],
+            ['./images/earth_1.jpg', './images/earth_2.jpg', './images/earth_3.jpg'],
+            ['./images/ydt_2.jpg', './images/ydt_1.jpg', './images/ydt_3.jpg'],
+           ]
+
+
 demo_mask = gr.Interface(fn=inference_mask1, 
                    inputs=[gr.ImageMask(brush_radius=8, label="prompt (提示图)"), gr.Image(label="img1 (测试图1)"), gr.Image(label="img2 (测试图2)")], 
                     #outputs=[gr.Image(shape=(448, 448), label="output1 (输出图1)"), gr.Image(shape=(448, 448), label="output2 (输出图2)")],
@@ -77,6 +110,29 @@ demo_mask = gr.Interface(fn=inference_mask1,
                    )
 
 
+
+demo_mask_sam = gr.Interface(fn=inference_mask1_sam,
+                   inputs=[gr.ImageMask(brush_radius=4, label="prompt (提示图)"), gr.Image(label="img1 (测试图1)"), gr.Image(label="img2 (测试图2)")],
+                    #outputs=[gr.Image(shape=(448, 448), label="output1 (输出图1)"), gr.Image(shape=(448, 448), label="output2 (输出图2)")],
+                    # outputs=[gr.Image(label="output1 (输出图1)").style(height=256, width=256), gr.Image(label="output2 (输出图2)").style(height=256, width=256)],
+                    #outputs=gr.Gallery(label="outputs (输出图)"),
+                    outputs=[gr.Image(label="SAM output (mask)").style(height=256, width=256),gr.Image(label="output1 (输出图1)").style(height=256, width=256), gr.Image(label="output2 (输出图2)").style(height=256, width=256)],
+                    # outputs=[gr.Image(label="output3 (输出图1)").style(height=256, width=256), gr.Image(label="output4 (输出图2)").style(height=256, width=256)],
+                    examples=examples_sam,
+                    #title="SegGPT for Any Segmentation<br>(Painter Inside)",
+                    description="<p> \
+                    Choose an example below &#128293; &#128293;  &#128293; <br>\
+                    Or, upload by yourself: <br>\
+                    1. Upload images to be tested to 'img1' and 'img2'. <br>2. Upload a prompt image to 'prompt' and draw <strong>a point or line on the target<strong>.  <br>\
+                            <br> \
+                            💎 SAM segment the target with any point or scribble, then SegGPT segments all other images. <br>\
+                            💎 Examples below were never trained and are randomly selected for testing in the wild. <br>\
+                            💎 Current UI interface only unleashes a small part of the capabilities of SegGPT, i.e., 1-shot case. \
+</p>",
+                   cache_examples=False,
+                   allow_flagging="never",
+                   )
+
 title = "SegGPT: Segmenting Everything In Context<br> \
 <div align='center'> \
 <h2><a href='https://arxiv.org/abs/2304.03284' target='_blank' rel='noopener'>[paper]</a> \
@@ -87,7 +143,7 @@ title = "SegGPT: Segmenting Everything In Context<br> \
 </div> \
 "
 
-demo = gr.TabbedInterface([demo_mask, ], ['General 1-shot', ], title=title)
+demo = gr.TabbedInterface([demo_mask_sam, demo_mask], ['SAM+SegGPT (一触百通)', 'General 1-shot'], title=title)
 
 #demo.launch(share=True, auth=("baai", "vision"))
 demo.launch(enable_queue=False)