Spaces:
Running
Running
| import cv2 | |
| import gradio as gr | |
| import numpy as np | |
| import onnxruntime | |
| from scipy.ndimage.filters import gaussian_filter | |
| from skimage.color import rgb2gray | |
| def xdog(im, gamma=0.98, phi=200, eps=-0.1, k=1.6, sigma=1): | |
| # Source : https://github.com/CemalUnal/XDoG-Filter | |
| # Reference : XDoG: An eXtended difference-of-Gaussians compendium including advanced image stylization | |
| # Link : http://citeseerx.ist.psu.edu/viewdoc/download?doi=10.1.1.365.151&rep=rep1&type=pdf | |
| if im.shape[2] == 3: | |
| im = rgb2gray(im) | |
| imf1 = gaussian_filter(im, sigma) | |
| imf2 = gaussian_filter(im, sigma * k) | |
| imdiff = imf1 - gamma * imf2 | |
| imdiff = (imdiff < eps) * 1.0 + (imdiff >= eps) * (1.0 + np.tanh(phi * imdiff)) | |
| imdiff -= imdiff.min() | |
| imdiff /= imdiff.max() | |
| imdiff *= 255.0 | |
| imdiff = imdiff.astype('uint8') | |
| imdiff = cv2.cvtColor(imdiff, cv2.COLOR_GRAY2BGR) | |
| return (imdiff / 255.0).astype('float32') | |
| def swin_model(img): | |
| h, w = img.shape[0], img.shape[1] | |
| factor = max(h / 512, w / 512) | |
| if factor > 1: | |
| img = cv2.resize(img, (int(h / factor), int(w / factor))) | |
| img = np.transpose(img[:, :, [2, 1, 0]], (2, 0, 1)) / 255.0 | |
| img = img[None, ...].astype(np.float32) | |
| mod_pad_h, mod_pad_w = 0, 0 | |
| window_size = 8 | |
| if h % window_size != 0: | |
| mod_pad_h = window_size - h % window_size | |
| if w % window_size != 0: | |
| mod_pad_w = window_size - w % window_size | |
| img = np.pad(img, (mod_pad_w, mod_pad_h), 'reflect') | |
| ort_session = onnxruntime.InferenceSession('model.onnx', providers=['CPUExecutionProvider']) | |
| ort_inputs = {ort_session.get_inputs()[0].name: img} | |
| output = ort_session.run(None, ort_inputs)[0] | |
| _, _, a, b = output.shape | |
| output = output[:, :, 0:a - mod_pad_h, 0:b - mod_pad_w] | |
| output = np.squeeze(output).clip(0, 1) | |
| if factor > 1: | |
| output = cv2.resize(output, h, w) | |
| return output | |
| demo = gr.Blocks() | |
| with demo: | |
| gr.Markdown( | |
| """ | |
| # Hello,World! | |
| This is an sketch extraction using in anime art by swin_transformer, you can see the difference between X-dog and learning-based algorithm. | |
| """) | |
| input_image = gr.Image() | |
| output_xdog = gr.Image() | |
| output_deep = gr.Image() | |
| b_xdog = gr.Button("X_dog") | |
| b_learning_based = gr.Button("learning based") | |
| b_xdog.click(xdog, inputs=input_image, outputs=output_xdog) | |
| b_learning_based.click(swin_model, inputs=input_image, outputs=output_deep) | |
| demo.launch() | |