Upload app.py

app.py

@@ -80,6 +80,86 @@ def check_input_image(input_image):
 
 
 def preprocess(input_image, do_remove_background, foreground_ratio):
+    def pre_process(img: np.array) -> np.array:
+        # H, W, C -> C, H, W
+        img = np.transpose(img[:, :, 0:3], (2, 0, 1))
+        # C, H, W -> 1, C, H, W
+        img = np.expand_dims(img, axis=0).astype(np.float32)
+        return img
+
+
+    def post_process(img: np.array) -> np.array:
+        # 1, C, H, W -> C, H, W
+        img = np.squeeze(img)
+        # C, H, W -> H, W, C
+        img = np.transpose(img, (1, 2, 0))[:, :, ::-1].astype(np.uint8)
+        return img
+
+
+    def inference(model_path: str, img_array: np.array) -> np.array:
+        options = onnxruntime.SessionOptions()
+        options.intra_op_num_threads = 1
+        options.inter_op_num_threads = 1
+        ort_session = onnxruntime.InferenceSession(model_path, options)
+        ort_inputs = {ort_session.get_inputs()[0].name: img_array}
+        ort_outs = ort_session.run(None, ort_inputs)
+
+        return ort_outs[0]
+
+
+    def convert_pil_to_cv2(input_image):
+        # pil_image = image.convert("RGB")
+        open_cv_image = np.array(input_image)
+        # RGB to BGR
+        open_cv_image = open_cv_image[:, :, ::-1].copy()
+        return open_cv_image
+
+
+    def upscale(image, model):
+        model_path = f"models/{model}.ort"
+        img = convert_pil_to_cv2(image)
+        if img.ndim == 2:
+            img = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)
+
+        if img.shape[2] == 4:
+            alpha = img[:, :, 3]  # GRAY
+            alpha = cv2.cvtColor(alpha, cv2.COLOR_GRAY2BGR)  # BGR
+            alpha_output = post_process(inference(model_path, pre_process(alpha)))  # BGR
+            alpha_output = cv2.cvtColor(alpha_output, cv2.COLOR_BGR2GRAY)  # GRAY
+
+            img = img[:, :, 0:3]  # BGR
+            image_output = post_process(inference(model_path, pre_process(img)))  # BGR
+            image_output = cv2.cvtColor(image_output, cv2.COLOR_BGR2BGRA)  # BGRA
+            image_output[:, :, 3] = alpha_output
+
+        elif img.shape[2] == 3:
+            image_output = post_process(inference(model_path, pre_process(img)))  # BGR
+
+        return image_output
+
+
+
+    def fill_background(image):
+        image = np.array(image).astype(np.float32) / 255.0
+        image = image[:, :, :3] * image[:, :, 3:4] + (1 - image[:, :, 3:4]) * 0.5
+        image = Image.fromarray((image * 255.0).astype(np.uint8))
+        return image
+
+
+
+
+    if do_remove_background:
+        image = image_output.convert("RGB")
+        image = remove_background(image, rembg_session)
+        image = resize_foreground(image, foreground_ratio)
+        image = fill_background(image)
+    else:
+        image = image_output
+        if image.mode == "RGBA":
+            image = fill_background(image)
+    return image
+
+    
     def fill_background(image):
         image = np.array(image).astype(np.float32) / 255.0
         image = image[:, :, :3] * image[:, :, 3:4] + (1 - image[:, :, 3:4]) * 0.5