bryandlee / animegan2-pytorch

PyTorch implementation of AnimeGANv2

Geek Repo:Geek Repo

Github PK Tool:Github PK Tool

PNG images with transparent background are not supported

njfanxun opened this issue · comments

The portrait is on a png image with a transparent background and cannot be stylized。How to modify the code to support it?

Hi, can I see the error message?

no error,but the generated image will automatically fill the background。You can try it out with a png with an alpha channel in the background. I think maybe the alpha value is discarded when the image data is processed. But I don't know how to modify the code. code reference demo.ipynb。

1
anime_full_1

Here you go, replace the last cell in demo.ipynb with this:

import requests
from torchvision.transforms.functional import to_tensor, to_pil_image

# Load image
image = Image.open(requests.get("https://user-images.githubusercontent.com/1555047/155689070-007698e6-eef3-497a-8aec-b26ffe1eb2dc.png", stream=True).raw)
image_rgb = image.convert("RGB")

# Detect faces & resize image s.t. face size = 256
face_detector = get_dlib_face_detector()
landmarks = face_detector(image_rgb)
landmarks = np.array(landmarks)
mean_face_size = (landmarks.max(axis=1) - landmarks.min(axis=1)).mean()

resize_ratio = 256 / max(mean_face_size, 32)
image_inference_size = (np.array(image.size) * resize_ratio).astype(int)
image_inference_size -= image_inference_size % 4
print(image_inference_size)
image_rgb = image_rgb.resize(image_inference_size, Image.LANCZOS)

# Model inference
model_input = to_tensor(image_rgb).unsqueeze(0) * 2 - 1
with torch.inference_mode():
    image_output = model(model_input.to(device))
image_output = (image_output * 0.5 + 0.5).clip(0, 1)
image_output = to_pil_image(image_output.cpu()[0])

# Resize image back to the original size & mode
image_output = image_output.resize(image.size, Image.LANCZOS).convert(image.mode)
if image.mode == "RGBA":  # <<<<< THIS PART HANDLES WHAT YOU WANTED
    alpha = image.split()[-1]
    image_output.putalpha(alpha)

display(image_output)

The output will look like this:

image

+) image size and face ratio does matter, you can get cleaner results if you use aligned & cropped faces

image

oh awesome thank you so much