Add background color

app.py

@@ -5,7 +5,7 @@ from modelscope.utils.constant import Tasks
 
 portrait_matting = pipeline(Tasks.portrait_matting, model='damo/cv_unet_image-matting')
 
-def matting(image):
+def matting(image, use_color, color):
     result = portrait_matting(image)
     image_pil = Image.fromarray(result['output_img'])
     alpha_channel = image_pil.split()[3]
@@ -15,36 +15,42 @@ def matting(image):
     mask_image.paste(alpha_channel, alpha_channel)
 
     # new image with background
-    new_image = Image.new("RGBA", image_pil.size, (207,195,181))
-    new_image.paste(image_pil, (0,0), image_pil)
-    new_image = new_image.convert("RGB")
+    if use_color:
+        new_image = Image.new("RGBA", image_pil.size, color)
+        new_image.paste(image_pil, (0,0), image_pil)
+        new_image = new_image.convert("RGB")
+    else:
+        new_image = image_pil
 
-    return [mask_image, new_image]
+    return [mask_image, new_image] 
 
 def merge(org_image, add_image):
     add_image = add_image.resize(org_image.size)
     org_image.paste(add_image, (0,0), add_image)
-    return org_image
+    return [org_image]
 
 with gr.Blocks() as demo:
     with gr.Tab(label="Portrait Matting"):
         with gr.Row():
             with gr.Column():
-                image = gr.Image()
+                image = gr.Image(height="40vh")
+                with gr.Row():
+                    use_color = gr.Checkbox(label="use background color", value=True, container=False)
+                    color = gr.ColorPicker(info="background color",value="#89cff0", container=False)
                 btn_matt= gr.Button()
-            result_matt = gr.Gallery()
+            result_matt = gr.Gallery(height="40vh")
     with gr.Tab(label="Merge"):
         with gr.Row():
             with gr.Column():
                 with gr.Row():
-                    org_image = gr.Image(label="background", type='pil', image_mode='RGBA')
-                    add_image = gr.Image(label="foreground", type='pil', image_mode='RGBA')
+                    org_image = gr.Image(label="background", type='pil', image_mode='RGBA', height="40vh") 
+                    add_image = gr.Image(label="foreground", type='pil', image_mode='RGBA', height="40vh")
                 btn_merge = gr.Button()
-            result_merge = gr.Image()
-
+            result_merge = gr.Gallery(height="40vh")
+    
     btn_matt.click(
         fn=matting,
-        inputs=[image],
+        inputs=[image, use_color, color],
         outputs=[result_matt],
     )
 
@@ -54,4 +60,4 @@ with gr.Blocks() as demo:
         outputs=[result_merge],
     )
 
-demo.launch()
+demo.launch()