RaiNote/ComicStyleTransferU2Net

1# https://huggingface.co/spaces/Norod78/ComicsHeroU2Net/blob/main/app.py as template/basics2import os3from pathlib import Path4import subprocess5import sys6import PIL7import cv2 as cv8import torch9from PIL import Image10import gradio as gr11import numpy as np12 13torch.set_grad_enabled(False)14 15 16# https://en.wikipedia.org/wiki/Unsharp_masking17# https://stackoverflow.com/a/55590133/149560618def unsharp_mask(image, kernel_size=(5, 5), sigma=1.0, amount=2.0, threshold=0):19    """Return a sharpened version of the image, using an unsharp mask."""20    blurred = cv.GaussianBlur(image, kernel_size, sigma)21    sharpened = float(amount + 1) * image - float(amount) * blurred22    sharpened = np.maximum(sharpened, np.zeros(sharpened.shape))23    sharpened = np.minimum(sharpened, 255 * np.ones(sharpened.shape))24    sharpened = sharpened.round()25    if threshold > 0:26        low_contrast_mask = np.absolute(image - blurred) < threshold27        np.copyto(sharpened, image, where=low_contrast_mask)28    return sharpened29 30 31def normPRED(d):32    ma = np.max(d)33    mi = np.min(d)34 35    dn = (d - mi) / (ma - mi)36 37    return dn38 39 40def array_to_np(array_in):41    array_in = normPRED(array_in)42    array_in = np.squeeze(255.0 * array_in)43    array_in = np.transpose(array_in, (1, 2, 0))44    return array_in45 46 47def image_as_array(image_in):48    image_in = np.array(image_in, np.float32)49    tmpImg = np.zeros((image_in.shape[0], image_in.shape[1], 3))50    image_in = image_in / np.max(image_in)51    if image_in.shape[2] == 1:52        tmpImg[:, :, 0] = (image_in[:, :, 0] - 0.485) / 0.22953        tmpImg[:, :, 1] = (image_in[:, :, 0] - 0.485) / 0.22954        tmpImg[:, :, 2] = (image_in[:, :, 0] - 0.485) / 0.22955    else:56        tmpImg[:, :, 0] = (image_in[:, :, 0] - 0.485) / 0.22957        tmpImg[:, :, 1] = (image_in[:, :, 1] - 0.456) / 0.22458        tmpImg[:, :, 2] = (image_in[:, :, 2] - 0.406) / 0.22559 60    tmpImg = tmpImg.transpose((2, 0, 1))61    image_out = np.expand_dims(tmpImg, 0)62    return image_out63 64 65def face2hero(model,66              img: Image.Image,67              size: int,68              unsharpen: bool = False69              ) -> Image.Image:70    temp_img = img.resize((size, size))71    input = torch.Tensor(image_as_array(temp_img))72    results = model(input)73    hero_np_image = array_to_np(results[1].detach().numpy())74    if unsharpen:75        hero_np_image = unsharp_mask(hero_np_image)76    hero_image = Image.fromarray(hero_np_image.astype(np.uint8))77 78    # output = img_concat_h(array_to_image(aligned_img), hero_image)79    del results80 81    return hero_image82 83 84def inference(img: Image.Image, unsharpen, model_path):85    model = torch.jit.load(model_path)86    model.eval()87    out = face2hero(model, img, 320, unsharpen)88    out = out.resize((img.width, img.height))89    return out90 91 92title = "Comic StyleTransfer U2Net"93description = "Style transfer a face into one of a \"Comics Style Transfer\""94article = ""95examples = []96print(torch.cuda.is_available())97demo = gr.Interface(98    inference,99    inputs=[gr.Image(type="pil", label="Input"), gr.Checkbox(label="Enable Unsharpen"),100            gr.Dropdown(show_label=True, label="Select Model", choices=list(Path(".").glob(f"*.jit.pt")),101                        max_choices=1)],102    outputs=[gr.Image(type="pil", label="Output")],103    title=title,104    description=description,105    article=article,106    examples=examples,107    allow_flagging="never"108)109 110demo.queue()111demo.launch()112