initial commit

app.py

@@ -20,7 +20,11 @@ def post_process(img: np.array) -> np.array:
     return img
 
 
-def inference(img_array: np.array) -> np.array:
+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)
 
@@ -28,6 +32,7 @@ def inference(img_array: np.array) -> np.array:
 
 
 def convert_pil_to_cv2(image):
+    # pil_image = image.convert("RGB")
     open_cv_image = np.array(image)
     # RGB to BGR
     open_cv_image = open_cv_image[:, :, ::-1].copy()
@@ -35,6 +40,7 @@ def convert_pil_to_cv2(image):
 
 
 def upscale(image):
+    model_path = "models/model.ort"
     img = convert_pil_to_cv2(image)
     if img.ndim == 2:
         img = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)
@@ -42,26 +48,20 @@ def upscale(image):
     if img.shape[2] == 4:
         alpha = img[:, :, 3]  # GRAY
         alpha = cv2.cvtColor(alpha, cv2.COLOR_GRAY2BGR)  # BGR
-        alpha_output = post_process(inference(pre_process(alpha)))  # 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(pre_process(img)))  # 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(pre_process(img)))  # BGR
+        image_output = post_process(inference(model_path, pre_process(img)))  # BGR
 
     return image_output
 
 
-model_path = "models/model.ort"
-options = onnxruntime.SessionOptions()
-options.intra_op_num_threads = 1
-options.inter_op_num_threads = 1
-ort_session = onnxruntime.InferenceSession(model_path, options)
-
 examples = [f"examples/example_{i+1}.png" for i in range(5)]
 css = ".output-image, .input-image, .image-preview {height: 480px !important} "