import numpy as np import cv2 import onnxruntime import gradio as gr 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: ort_session = onnxruntime.InferenceSession(model_path) 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(image): # pil_image = image.convert("RGB") open_cv_image = np.array(image) # RGB to BGR open_cv_image = open_cv_image[:, :, ::-1].copy() return open_cv_image def main(image): model_path = "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 gr.Interface( main, gr.inputs.Image(type="pil"), "image", title="Image Upscaling 🦆", allow_flagging="never", css=".output-image, .input-image, .image-preview {height: 500px !important} ", ).launch()