File size: 2,525 Bytes
a3d6c18 |
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 |
"""
Source url: https://github.com/OPHoperHPO/image-background-remove-tool
Author: Nikita Selin (OPHoperHPO)[https://github.com/OPHoperHPO].
License: Apache License 2.0
"""
import PIL.Image
import torch
from carvekit.utils.image_utils import to_tensor
def composite(
foreground: PIL.Image.Image,
background: PIL.Image.Image,
alpha: PIL.Image.Image,
device="cpu",
):
"""
Composites foreground with background by following
https://pymatting.github.io/intro.html#alpha-matting math formula.
Args:
device: Processing device
foreground: Image that will be pasted to background image with following alpha mask.
background: Background image
alpha: Alpha Image
Returns:
Composited image as PIL.Image instance.
"""
foreground = foreground.convert("RGBA")
background = background.convert("RGBA")
alpha_rgba = alpha.convert("RGBA")
alpha_l = alpha.convert("L")
fg = to_tensor(foreground).to(device)
alpha_rgba = to_tensor(alpha_rgba).to(device)
alpha_l = to_tensor(alpha_l).to(device)
bg = to_tensor(background).to(device)
alpha_l = alpha_l / 255
alpha_rgba = alpha_rgba / 255
bg = torch.where(torch.logical_not(alpha_rgba >= 1), bg, fg)
bg[:, :, 0] = alpha_l[:, :] * fg[:, :, 0] + (1 - alpha_l[:, :]) * bg[:, :, 0]
bg[:, :, 1] = alpha_l[:, :] * fg[:, :, 1] + (1 - alpha_l[:, :]) * bg[:, :, 1]
bg[:, :, 2] = alpha_l[:, :] * fg[:, :, 2] + (1 - alpha_l[:, :]) * bg[:, :, 2]
bg[:, :, 3] = alpha_l[:, :] * 255
del alpha_l, alpha_rgba, fg
return PIL.Image.fromarray(bg.cpu().numpy()).convert("RGBA")
def apply_mask(
image: PIL.Image.Image, mask: PIL.Image.Image, device="cpu"
) -> PIL.Image.Image:
"""
Applies mask to foreground.
Args:
device: Processing device.
image: Image with background.
mask: Alpha Channel mask for this image.
Returns:
Image without background, where mask was black.
"""
background = PIL.Image.new("RGBA", image.size, color=(130, 130, 130, 0))
return composite(image, background, mask, device=device).convert("RGBA")
def extract_alpha_channel(image: PIL.Image.Image) -> PIL.Image.Image:
"""
Extracts alpha channel from the RGBA image.
Args:
image: RGBA PIL image
Returns:
RGBA alpha channel image
"""
alpha = image.split()[-1]
bg = PIL.Image.new("RGBA", image.size, (0, 0, 0, 255))
bg.paste(alpha, mask=alpha)
return bg.convert("RGBA")
|