Duplicate from xiongjie/u2net_rgba

Co-authored-by: xiong jie <[email protected]>

.gitattributes

@@ -0,0 +1,27 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bin.* filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zstandard filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,38 @@
+---
+title: U2net_rgba
+emoji: 📉
+colorFrom: yellow
+colorTo: yellow
+sdk: gradio
+app_file: app.py
+pinned: false
+duplicated_from: xiongjie/u2net_rgba
+---
+
+# Configuration
+
+`title`: _string_  
+Display title for the Space
+
+`emoji`: _string_  
+Space emoji (emoji-only character allowed)
+
+`colorFrom`: _string_  
+Color for Thumbnail gradient (red, yellow, green, blue, indigo, purple, pink, gray)
+
+`colorTo`: _string_  
+Color for Thumbnail gradient (red, yellow, green, blue, indigo, purple, pink, gray)
+
+`sdk`: _string_  
+Can be either `gradio` or `streamlit`
+
+`sdk_version` : _string_  
+Only applicable for `streamlit` SDK.  
+See [doc](https://hf.co/docs/hub/spaces) for more info on supported versions.
+
+`app_file`: _string_  
+Path to your main application file (which contains either `gradio` or `streamlit` Python code).  
+Path is relative to the root of the repository.
+
+`pinned`: _boolean_  
+Whether the Space stays on top of your list.

app.py

@@ -0,0 +1,66 @@
+import os
+import copy
+import time
+
+import cv2 as cv
+import numpy as np
+import onnxruntime
+
+from PIL import Image
+
+import gradio
+
+def run_inference(onnx_session, input_size, image):
+    # リサイズ
+    temp_image = copy.deepcopy(image)
+    resize_image = cv.resize(temp_image, dsize=(input_size, input_size))
+    x = cv.cvtColor(resize_image, cv.COLOR_BGR2RGB)
+
+    # 前処理
+    x = np.array(x, dtype=np.float32)
+    mean = [0.485, 0.456, 0.406]
+    std = [0.229, 0.224, 0.225]
+    x = (x / 255 - mean) / std
+    x = x.transpose(2, 0, 1).astype('float32')
+    x = x.reshape(-1, 3, input_size, input_size)
+
+    # 推論
+    input_name = onnx_session.get_inputs()[0].name
+    output_name = onnx_session.get_outputs()[0].name
+    onnx_result = onnx_session.run([output_name], {input_name: x})
+
+    # 後処理
+    onnx_result = np.array(onnx_result).squeeze()
+    min_value = np.min(onnx_result)
+    max_value = np.max(onnx_result)
+    onnx_result = (onnx_result - min_value) / (max_value - min_value)
+    onnx_result *= 255
+    onnx_result = onnx_result.astype('uint8')
+
+    return onnx_result
+
+# Load model
+onnx_session = onnxruntime.InferenceSession("u2net.onnx")
+
+def create_rgba(mode, image):
+    out = run_inference(
+        onnx_session,
+        320,
+        image,
+    )
+    resize_image = cv.resize(out, dsize=(image.shape[1], image.shape[0]))
+
+    if mode == "binary":
+        resize_image[resize_image > 255] = 255
+        resize_image[resize_image < 125] = 0
+
+    mask = Image.fromarray(resize_image)
+
+    rgba_image = Image.fromarray(image).convert('RGBA')
+    rgba_image.putalpha(mask)
+
+    return rgba_image
+
+inputs = [gradio.inputs.Radio(["binary", "smooth"]), gradio.inputs.Image()]
+outputs = gradio.outputs.Image()
+gradio.Interface(fn=create_rgba, inputs=inputs, outputs=outputs).launch()

model_float32.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b069656b37769de408b6c47d5bc1d0d4a9a059925e3bc8d000a3c536026ca5cf
+size 4652116

packages.txt

@@ -0,0 +1 @@
+libgl1-mesa-dev

requirements.txt

@@ -0,0 +1,5 @@
+gradio
+numpy
+opencv-python
+onnxruntime==1.6.0
+Pillow

u2net.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:309ea1fc458e8f6711efb645da85d1cc2a9cd2aec261d379bba31c0a0ddc78af
+size 175995934