Spaces:
Running
on
T4
Running
on
T4
File size: 5,619 Bytes
1eaf723 659d1e7 dbe4fe6 6063bf9 659d1e7 dbe4fe6 cb897c1 dbe4fe6 cb897c1 dbe4fe6 feaa03b dbe4fe6 cb897c1 659d1e7 3c69e99 fa3cf6d 34974b2 2fb1494 34974b2 2fb1494 3c69e99 fa3cf6d 2fb1494 3c69e99 dbe4fe6 f8f84a0 c1307a8 dbe4fe6 48b9ade fa3cf6d dbe4fe6 d426447 228b8bd c36100a |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 |
import os
from huggingface_hub import hf_hub_download
os.system("pip -qq install facenet_pytorch")
from facenet_pytorch import MTCNN
from torchvision import transforms
import torch, PIL
from tqdm.notebook import tqdm
import gradio as gr
import torch
modelarcanev4 = hf_hub_download(repo_id="akhaliq/ArcaneGANv0.4", filename="ArcaneGANv0.4.jit")
modelarcanev3 = hf_hub_download(repo_id="akhaliq/ArcaneGANv0.3", filename="ArcaneGANv0.3.jit")
modelarcanev2 = hf_hub_download(repo_id="akhaliq/ArcaneGANv0.2", filename="ArcaneGANv0.2.jit")
mtcnn = MTCNN(image_size=256, margin=80)
# simplest ye olde trustworthy MTCNN for face detection with landmarks
def detect(img):
# Detect faces
batch_boxes, batch_probs, batch_points = mtcnn.detect(img, landmarks=True)
# Select faces
if not mtcnn.keep_all:
batch_boxes, batch_probs, batch_points = mtcnn.select_boxes(
batch_boxes, batch_probs, batch_points, img, method=mtcnn.selection_method
)
return batch_boxes, batch_points
# my version of isOdd, should make a separate repo for it :D
def makeEven(_x):
return _x if (_x % 2 == 0) else _x+1
# the actual scaler function
def scale(boxes, _img, max_res=1_500_000, target_face=256, fixed_ratio=0, max_upscale=2, VERBOSE=False):
x, y = _img.size
ratio = 2 #initial ratio
#scale to desired face size
if (boxes is not None):
if len(boxes)>0:
ratio = target_face/max(boxes[0][2:]-boxes[0][:2]);
ratio = min(ratio, max_upscale)
if VERBOSE: print('up by', ratio)
if fixed_ratio>0:
if VERBOSE: print('fixed ratio')
ratio = fixed_ratio
x*=ratio
y*=ratio
#downscale to fit into max res
res = x*y
if res > max_res:
ratio = pow(res/max_res,1/2);
if VERBOSE: print(ratio)
x=int(x/ratio)
y=int(y/ratio)
#make dimensions even, because usually NNs fail on uneven dimensions due skip connection size mismatch
x = makeEven(int(x))
y = makeEven(int(y))
size = (x, y)
return _img.resize(size)
"""
A useful scaler algorithm, based on face detection.
Takes PIL.Image, returns a uniformly scaled PIL.Image
boxes: a list of detected bboxes
_img: PIL.Image
max_res: maximum pixel area to fit into. Use to stay below the VRAM limits of your GPU.
target_face: desired face size. Upscale or downscale the whole image to fit the detected face into that dimension.
fixed_ratio: fixed scale. Ignores the face size, but doesn't ignore the max_res limit.
max_upscale: maximum upscale ratio. Prevents from scaling images with tiny faces to a blurry mess.
"""
def scale_by_face_size(_img, max_res=1_500_000, target_face=256, fix_ratio=0, max_upscale=2, VERBOSE=False):
boxes = None
boxes, _ = detect(_img)
if VERBOSE: print('boxes',boxes)
img_resized = scale(boxes, _img, max_res, target_face, fix_ratio, max_upscale, VERBOSE)
return img_resized
size = 256
means = [0.485, 0.456, 0.406]
stds = [0.229, 0.224, 0.225]
t_stds = torch.tensor(stds).cuda().half()[:,None,None]
t_means = torch.tensor(means).cuda().half()[:,None,None]
def makeEven(_x):
return int(_x) if (_x % 2 == 0) else int(_x+1)
img_transforms = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize(means,stds)])
def tensor2im(var):
return var.mul(t_stds).add(t_means).mul(255.).clamp(0,255).permute(1,2,0)
def proc_pil_img(input_image, model):
transformed_image = img_transforms(input_image)[None,...].cuda().half()
with torch.no_grad():
result_image = model(transformed_image)[0]
output_image = tensor2im(result_image)
output_image = output_image.detach().cpu().numpy().astype('uint8')
output_image = PIL.Image.fromarray(output_image)
return output_image
modelv4 = torch.jit.load(modelarcanev4).eval().cuda().half()
modelv3 = torch.jit.load(modelarcanev3).eval().cuda().half()
modelv2 = torch.jit.load(modelarcanev2).eval().cuda().half()
def process(im, version):
if version == 'version 0.4':
im = scale_by_face_size(im, target_face=256, max_res=1_500_000, max_upscale=1)
res = proc_pil_img(im, modelv4)
elif version == 'version 0.3':
im = scale_by_face_size(im, target_face=256, max_res=1_500_000, max_upscale=1)
res = proc_pil_img(im, modelv3)
else:
im = scale_by_face_size(im, target_face=256, max_res=1_500_000, max_upscale=1)
res = proc_pil_img(im, modelv2)
return res
title = "ArcaneGAN"
description = "Gradio demo for ArcaneGAN, portrait to Arcane style. To use it, simply upload your image, or click one of the examples to load them. Read more at the links below."
article = "<div style='text-align: center;'>ArcaneGan by <a href='https://twitter.com/devdef' target='_blank'>Alexander S</a> | <a href='https://github.com/Sxela/ArcaneGAN' target='_blank'>Github Repo</a> | <center><img src='https://visitor-badge.glitch.me/badge?page_id=akhaliq_arcanegan' alt='visitor badge'></center></div>"
gr.Interface(
process,
[gr.inputs.Image(type="pil", label="Input"),gr.inputs.Radio(choices=['version 0.2','version 0.3','version 0.4'], type="value", default='version 0.4', label='version')
],
gr.outputs.Image(type="pil", label="Output"),
title=title,
description=description,
article=article,
examples=[['bill.png','version 0.3'],['keanu.png','version 0.4'],['will.jpeg','version 0.4']],
allow_flagging=False,
allow_screenshot=False
).launch()
|