Playing with VitMatte for image matting interface
Tons of Application in:
- Movie special effects
- Video conference
- Image Editing
- Digital person creation
The idea is to separate a foreground object and background object.
Based formula I = alpha * F + (1-alpha) * B
Install Libraries
!pip install -q git+https://github.com/huggingface/transformers.git
Load image an d trimap
import matplotlib.pyplot as plt
from PIL import Image
import requests
url = "https://github.com/hustvl/ViTMatte/blob/main/demo/bulb_rgb.png?raw=true"
image = Image.open(requests.get(url, stream=True).raw).convert("RGB")
url = "https://github.com/hustvl/ViTMatte/blob/main/demo/bulb_trimap.png?raw=true"
trimap = Image.open(requests.get(url, stream=True).raw)
plt.figure(figsize=(15, 15))
plt.subplot(1, 2, 1)
plt.imshow(image)
plt.subplot(1, 2, 2)
plt.imshow(trimap)
plt.show()
# credit --> https://github.com/hustvl/ViTMatte
Visualize Foreground
import PIL
def cal_foreground(image: PIL.Image, alpha: PIL.Image):
"""
Calculate the foreground of the image.
Returns:
foreground: the foreground of the image, numpy array
"""
image = image.convert("RGB")
alpha = alpha.convert("L")
alpha = F.to_tensor(alpha).unsqueeze(0)
image = F.to_tensor(image).unsqueeze(0)
foreground = image * alpha + (1 - alpha)
foreground = foreground.squeeze(0).permute(1, 2, 0).numpy()
return foreground
fg = cal_foreground(image, prediction1)
plt.figure(figsize=(7, 7))
plt.imshow(fg)
plt.show()
# credit --> https://github.com/hustvl/ViTMatte
Background Replacement
image_path = "malibu.jpg"
background = Image.open(image_path)
plt.imshow(background)
def merge_new_bg(image, background, alpha):
"""
Merge the alpha matte with a new background.
Returns:
foreground: the foreground of the image, numpy array
"""
image = image.convert('RGB')
bg = background.convert('RGB')
alpha = alpha.convert('L')
image = F.to_tensor(image)
bg = F.to_tensor(bg)
bg = F.resize(bg, image.shape[-2:])
alpha = F.to_tensor(alpha)
new_image = image * alpha + bg * (1 - alpha)
new_image = new_image.squeeze(0).permute(1, 2, 0).numpy()
return new_image
# merge with new background
new_image = merge_new_bg(image, background, prediction)
plt.figure(figsize=(7, 7))
plt.imshow(new_image)
plt.show()