add face2paint
This commit is contained in:
parent
4a9c882a0a
commit
c78943efde
33
hubconf.py
33
hubconf.py
|
@ -1,6 +1,6 @@
|
|||
import torch
|
||||
|
||||
def animegan2(pretrained=True, device="cpu", progress=True, check_hash=True):
|
||||
def generator(pretrained=True, device="cpu", progress=True, check_hash=True):
|
||||
release_url = "https://github.com/xhlulu/animegan2-pytorch/releases/download/weights"
|
||||
known = {
|
||||
name: f"{release_url}/{name}.pt"
|
||||
|
@ -31,4 +31,33 @@ def animegan2(pretrained=True, device="cpu", progress=True, check_hash=True):
|
|||
)
|
||||
model.load_state_dict(state_dict)
|
||||
|
||||
return model
|
||||
return model
|
||||
|
||||
|
||||
def face2paint(device="cpu"):
|
||||
from PIL import Image
|
||||
from torchvision.transforms.functional import to_tensor, to_pil_image
|
||||
|
||||
def face2paint(
|
||||
model: torch.nn.Module,
|
||||
img: Image.Image,
|
||||
size: int,
|
||||
side_by_side: bool = True,
|
||||
device: str = device,
|
||||
) -> Image.Image:
|
||||
w, h = img.size
|
||||
s = min(w, h)
|
||||
img = img.crop(((w - s) // 2, (h - s) // 2, (w + s) // 2, (h + s) // 2))
|
||||
img = img.resize((size, size), Image.LANCZOS)
|
||||
|
||||
input = to_tensor(img).unsqueeze(0) * 2 - 1
|
||||
output = model(input.to(device)).cpu()[0]
|
||||
|
||||
if side_by_side:
|
||||
output = torch.cat([input[0], output], dim=2)
|
||||
|
||||
output = (output * 0.5 + 0.5).clip(0, 1)
|
||||
|
||||
return to_pil_image(output)
|
||||
|
||||
return face2paint
|
||||
|
|
Loading…
Reference in a new issue