add: 白背景モード実装

README.md

@@ -4,7 +4,7 @@ emoji: 🪄🖼️
 colorFrom: indigo
 colorTo: pink
 sdk: gradio
-sdk_version: 3.1.4
+sdk_version: 3.50.2
 app_file: app.py
 pinned: false
 license: apache-2.0

app.py

@@ -4,7 +4,6 @@ import onnxruntime as rt
 import numpy as np
 import cv2
 
-
 def get_mask(img, s=1024):
     img = (img / 255).astype(np.float32)
     h, w = h0, w0 = img.shape[:-1]
@@ -20,17 +19,18 @@ def get_mask(img, s=1024):
     mask = cv2.resize(mask, (w0, h0))[:, :, np.newaxis]
     return mask
 
-
-def rmbg_fn(img):
+def rmbg_fn(img, bg_color):
     mask = get_mask(img)
-    img = (mask * img + 255 * (1 - mask)).astype(np.uint8)
-    mask = (mask * 255).astype(np.uint8)
-    img = (mask * img + (1 - mask) * 255).astype(np.uint8)
-    # img = np.concatenate([img, mask], axis=2, dtype=np.uint8)
-    mask = (mask * 255).astype(np.uint8)
+    if bg_color == "透過":
+        img_with_bg = (mask * img + 255 * (1 - mask)).astype(np.uint8)
+        mask = (mask * 255).astype(np.uint8)
+        img_with_bg = np.concatenate([img_with_bg, mask], axis=2, dtype=np.uint8)
+    else:  # 白色背景
+        foreground = mask * img
+        background = 255 * (1 - mask)  # 白背景
+        img_with_bg = (foreground + background).astype(np.uint8)
     mask = mask.repeat(3, axis=2)
-    return mask, img
-
+    return mask, img_with_bg
 
 if __name__ == "__main__":
     providers = ['CUDAExecutionProvider', 'CPUExecutionProvider']
@@ -44,11 +44,12 @@ if __name__ == "__main__":
         with gr.Row():
             with gr.Column():
                 input_img = gr.Image(label="input image")
+                bg_color = gr.Radio(choices=["透過", "白色"], value="透過", label="背景色")
                 examples_data = [[f"examples/{x:02d}.jpg"] for x in range(1, 4)]
                 examples = gr.Dataset(components=[input_img], samples=examples_data)
-            run_btn = gr.Button(variant="primary")
+                run_btn = gr.Button(variant="primary")
             output_mask = gr.Image(label="mask")
             output_img = gr.Image(label="result", image_mode="RGBA")
         examples.click(lambda x: x[0], [examples], [input_img])
-        run_btn.click(rmbg_fn, [input_img], [output_mask, output_img])
-    app.launch()
+        run_btn.click(rmbg_fn, [input_img, bg_color], [output_mask, output_img])
+    app.launch(server_name='0.0.0.0')

requirements.txt

@@ -1,4 +1,5 @@
-onnx
-onnxruntime-gpu
-opencv-python
-numpy
+onnx==1.15.0
+onnxruntime-gpu==1.17.1
+opencv-python-headless==4.9.0.80
+numpy==1.26.3
+