Spaces:
Sleeping
Sleeping
| # https://huggingface.co/spaces/Norod78/ComicsHeroU2Net/blob/main/app.py as template/basics | |
| import os | |
| from pathlib import Path | |
| import subprocess | |
| import sys | |
| import PIL | |
| import cv2 as cv | |
| import torch | |
| from PIL import Image | |
| import gradio as gr | |
| import numpy as np | |
| torch.set_grad_enabled(False) | |
| # https://en.wikipedia.org/wiki/Unsharp_masking | |
| # https://stackoverflow.com/a/55590133/1495606 | |
| def unsharp_mask(image, kernel_size=(5, 5), sigma=1.0, amount=2.0, threshold=0): | |
| """Return a sharpened version of the image, using an unsharp mask.""" | |
| blurred = cv.GaussianBlur(image, kernel_size, sigma) | |
| sharpened = float(amount + 1) * image - float(amount) * blurred | |
| sharpened = np.maximum(sharpened, np.zeros(sharpened.shape)) | |
| sharpened = np.minimum(sharpened, 255 * np.ones(sharpened.shape)) | |
| sharpened = sharpened.round() | |
| if threshold > 0: | |
| low_contrast_mask = np.absolute(image - blurred) < threshold | |
| np.copyto(sharpened, image, where=low_contrast_mask) | |
| return sharpened | |
| def normPRED(d): | |
| ma = np.max(d) | |
| mi = np.min(d) | |
| dn = (d - mi) / (ma - mi) | |
| return dn | |
| def array_to_np(array_in): | |
| array_in = normPRED(array_in) | |
| array_in = np.squeeze(255.0 * array_in) | |
| array_in = np.transpose(array_in, (1, 2, 0)) | |
| return array_in | |
| def image_as_array(image_in): | |
| image_in = np.array(image_in, np.float32) | |
| tmpImg = np.zeros((image_in.shape[0], image_in.shape[1], 3)) | |
| image_in = image_in / np.max(image_in) | |
| if image_in.shape[2] == 1: | |
| tmpImg[:, :, 0] = (image_in[:, :, 0] - 0.485) / 0.229 | |
| tmpImg[:, :, 1] = (image_in[:, :, 0] - 0.485) / 0.229 | |
| tmpImg[:, :, 2] = (image_in[:, :, 0] - 0.485) / 0.229 | |
| else: | |
| tmpImg[:, :, 0] = (image_in[:, :, 0] - 0.485) / 0.229 | |
| tmpImg[:, :, 1] = (image_in[:, :, 1] - 0.456) / 0.224 | |
| tmpImg[:, :, 2] = (image_in[:, :, 2] - 0.406) / 0.225 | |
| tmpImg = tmpImg.transpose((2, 0, 1)) | |
| image_out = np.expand_dims(tmpImg, 0) | |
| return image_out | |
| def face2hero(model, | |
| img: Image.Image, | |
| size: int, | |
| unsharpen: bool = False | |
| ) -> Image.Image: | |
| temp_img = img.resize((size, size)) | |
| input = torch.Tensor(image_as_array(temp_img)) | |
| results = model(input) | |
| hero_np_image = array_to_np(results[1].detach().numpy()) | |
| if unsharpen: | |
| hero_np_image = unsharp_mask(hero_np_image) | |
| hero_image = Image.fromarray(hero_np_image.astype(np.uint8)) | |
| # output = img_concat_h(array_to_image(aligned_img), hero_image) | |
| del results | |
| return hero_image | |
| def inference(img: Image.Image, unsharpen, model_path): | |
| model = torch.jit.load(model_path) | |
| model.eval() | |
| out = face2hero(model, img, 320, unsharpen) | |
| out = out.resize((img.width, img.height)) | |
| return out | |
| title = "Comic StyleTransfer U2Net" | |
| description = "Style transfer a face into one of a \"Comics Style Transfer\"" | |
| article = "" | |
| examples = [] | |
| print(torch.cuda.is_available()) | |
| demo = gr.Interface( | |
| inference, | |
| inputs=[gr.Image(type="pil", label="Input"), gr.Checkbox(label="Enable Unsharpen"), | |
| gr.Dropdown(show_label=True, label="Select Model", choices=list(Path(".").glob(f"*.jit.pt")), | |
| max_choices=1)], | |
| outputs=[gr.Image(type="pil", label="Output")], | |
| title=title, | |
| description=description, | |
| article=article, | |
| examples=examples, | |
| allow_flagging="never" | |
| ) | |
| demo.queue() | |
| demo.launch() | |