| import os | |
| import cv2 | |
| import gradio as gr | |
| import numpy as np | |
| import onnxruntime as ort | |
| from PIL import Image | |
| _sess_options = ort.SessionOptions() | |
| _sess_options.intra_op_num_threads = os.cpu_count() | |
| MODEL_SESS = ort.InferenceSession( | |
| "cartoonizer.onnx", _sess_options, providers=["CPUExecutionProvider"] | |
| ) | |
| def preprocess_image(image: Image) -> np.ndarray: | |
| image = np.array(image) | |
| image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR) | |
| h, w, c = np.shape(image) | |
| if min(h, w) > 720: | |
| if h > w: | |
| h, w = int(720 * h / w), 720 | |
| else: | |
| h, w = 720, int(720 * w / h) | |
| image = cv2.resize(image, (w, h), interpolation=cv2.INTER_AREA) | |
| h, w = (h // 8) * 8, (w // 8) * 8 | |
| image = image[:h, :w, :] | |
| image = image.astype(np.float32) / 127.5 - 1 | |
| return np.expand_dims(image, axis=0) | |
| def inference(image: np.ndarray) -> Image: | |
| image = preprocess_image(image) | |
| results = MODEL_SESS.run(None, {"input_photo:0": image}) | |
| output = (np.squeeze(results[0]) + 1.0) * 127.5 | |
| output = np.clip(output, 0, 255).astype(np.uint8) | |
| output = cv2.cvtColor(output, cv2.COLOR_BGR2RGB) | |
| return Image.fromarray(output) | |
| title = "Generate cartoonized images" | |
| article = "Demo of CartoonGAN model (https://systemerrorwang.github.io/White-box-Cartoonization/). \nDemo image is from https://unsplash.com/photos/f0SgAs27BYI." | |
| iface = gr.Interface( | |
| inference, | |
| inputs=gr.inputs.Image(type="pil", label="Input Image"), | |
| outputs="image", | |
| title=title, | |
| article=article, | |
| allow_flagging="never", | |
| examples=[["mountain.jpeg"]], | |
| ) | |
| iface.launch() | |