Update app.py

app.py

@@ -2,34 +2,34 @@ import os
 import copy
 import time
 
-import cv2 as cv
 import numpy as np
 import onnxruntime
 
-from PIL import Image
+from PIL import Image, ImageOps
 
 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)
+    # Resize
+    temp_image = image.copy()
+    resize_image = temp_image.resize((input_size, input_size), Image.ANTIALIAS)
+    x = ImageOps.exif_transpose(resize_image)
+    x = np.array(x)
 
-    # 前処理
-    x = np.array(x, dtype=np.float32)
-    mean = [0.485, 0.456, 0.406]
-    std = [0.229, 0.224, 0.225]
+    # Preprocessing
+    x = x.astype(np.float32)
+    mean = np.array([0.485, 0.456, 0.406], dtype=np.float32)
+    std = np.array([0.229, 0.224, 0.225], dtype=np.float32)
     x = (x / 255 - mean) / std
     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})
 
-    # 後処理
+    # Postprocessing
     onnx_result = np.array(onnx_result).squeeze()
     min_value = np.min(onnx_result)
     max_value = np.max(onnx_result)
@@ -43,20 +43,20 @@ def run_inference(onnx_session, input_size, image):
 onnx_session = onnxruntime.InferenceSession("u2net.onnx")
 
 def create_rgba(mode, image):
+    image = Image.fromarray(image).convert('RGB')
     out = run_inference(
         onnx_session,
         320,
         image,
     )
-    resize_image = cv.resize(out, dsize=(image.shape[1], image.shape[0]))
+    resize_image = Image.fromarray(out).resize((image.size[0], image.size[1]), Image.ANTIALIAS)
 
     if mode == "binary":
-        resize_image[resize_image > 255] = 255
-        resize_image[resize_image < 125] = 0
+        resize_image = resize_image.point(lambda x: 255 if x > 125 else 0)
 
-    mask = Image.fromarray(resize_image)
+    mask = resize_image
 
-    rgba_image = Image.fromarray(image).convert('RGBA')
+    rgba_image = image.convert('RGBA')
     rgba_image.putalpha(mask)
 
     return rgba_image