Fix app.

app.py

@@ -11,42 +11,52 @@ from PIL import Image
 import gradio
 
 def run_inference(onnx_session, input_size, image):
-    # Pre process:Resize, BGR->RGB, Transpose, PyTorch standardization, float32 cast
+    # リサイズ
     temp_image = copy.deepcopy(image)
-    resize_image = cv.resize(temp_image, dsize=(input_size[0], input_size[1]))
+    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.reshape(-1, input_size[0], input_size[1], 3).astype('float32')
+    x = x.transpose(2, 0, 1).astype('float32')
+    x = x.reshape(-1, 3, input_size, input_size)
 
-    # Inference
+    # 推論
     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})
 
-    # Post process
+    # 後処理
     onnx_result = np.array(onnx_result).squeeze()
-    onnx_result = (1 - onnx_result)
     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')
+    onnx_result[onnx_result >= 125] = 255
+    onnx_result[onnx_result < 125] = 0
 
     return onnx_result
 
 # Load model
-onnx_session = onnxruntime.InferenceSession("model_float32.onnx")
+onnx_session = onnxruntime.InferenceSession("u2net.onnx")
 
 def create_rgba(image):
-  print(images.shape)
-  return run_inference(
-      onnx_session,
-      [480,640],
-      image,
-  )
+    out = run_inference(
+        onnx_session,
+        320,
+        image,
+    )
+    resize_image = cv.resize(out, dsize=(image.shape[1], image.shape[0]))
+    mask = Image.fromarray(resize_image)
+
+    rgba_image = Image.fromarray(image).convert('RGBA')
+    rgba_image.putalpha(mask)
+
+    return rgba_image
 
 css = ".output_image {height: 100% !important; width: 100% !important;}"
 inputs = gradio.inputs.Image()

u2net.onnx

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