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。
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 size and face ratio does matter, you can get cleaner results if you use aligned & cropped faces
oh awesome thank you so much