3-8. Background removal with RMBG v1.4 and OpenVINO

AI/OpenVINO

3-8. Background removal with RMBG v1.4 and OpenVINO

된다잘된다 2024. 7. 18. 14:28

참고: https://docs.openvino.ai/2024/notebooks/rmbg-background-removal-with-output.html

Background removal with RMBG v1.4 and OpenVINO — OpenVINO™ documentation

Background removal with RMBG v1.4 and OpenVINO This Jupyter notebook can be launched after a local installation only. Background matting is the process of accurately estimating the foreground object in images and videos. It is a very important technique in

docs.openvino.ai

이번 예제는 이미지 작업중에 흔히 말하는 누끼따기에 대한 내용이다.

영어로는 background matting이라고 한다는데, VFX같은 컴퓨터 그래픽이 들어가는 영상작업에 많이 사용된다.

갤럭시 광고에서도 사진에 펜으로 동그라미를 그리면 그 부분만 떼어냈던 것이 기억난다.

설명에는 이미지 구분(image segmentation)과는 다르다고 하는데,

background matting은 이미지의 픽셀들이 foreground에도 background에도 속할 수 있는 반면에 (partial pixels),

image segmentation은 픽셀마다 background인지 foreground인지 나누기 때문이라고 한다. (여기서는 안중요함)

사용하는 모델은 RMGB v1.4로 성능이 좋다고 한다. 더 자세한 사항은 이 링크를 참조하시길.

먼저 시작하기 전에 필요한 패키지들을 받아서 import하고 모델도 같이 다운 받기한다.

%pip install -q torch torchvision pillow huggingface_hub "openvino>=2024.0.0" matplotlib "gradio>=4.15" "transformers>=4.39.1" tqdm --extra-index-url https://download.pytorch.org/whl/cpu

나는 파일이 있는 위치를 모르는데 hf_hub_download()는 알고 있나보다.

이 모델은 PyTorch를 사용하므로 torch 모델을 load할 때 사용하는 method를 사용해야 한다. (지난 예시에서 했었다!)

from transformers import AutoModelForImageSegmentation

net = AutoModelForImageSegmentation.from_pretrained("briaai/RMBG-1.4", trust_remote_code=True)

일단 실행시켜서 잘 작동하는 지 부터 보자.

import torch
from PIL import Image
from utilities import preprocess_image, postprocess_image
import numpy as np
from matplotlib import pyplot as plt


def visualize_result(orig_img: Image, mask: Image, result_img: Image):
    """
    Helper for results visualization

    parameters:
       orig_img (Image): input image
       mask (Image): background mask
       result_img (Image) output image
    returns:
      plt.Figure: plot with 3 images for visualization
    """
    titles = ["Original", "Background Mask", "Without background"]
    im_w, im_h = orig_img.size
    is_horizontal = im_h <= im_w
    figsize = (20, 20)
    num_images = 3
    fig, axs = plt.subplots(
        num_images if is_horizontal else 1,
        1 if is_horizontal else num_images,
        figsize=figsize,
        sharex="all",
        sharey="all",
    )
    fig.patch.set_facecolor("white")
    list_axes = list(axs.flat)
    for a in list_axes:
        a.set_xticklabels([])
        a.set_yticklabels([])
        a.get_xaxis().set_visible(False)
        a.get_yaxis().set_visible(False)
        a.grid(False)
    list_axes[0].imshow(np.array(orig_img))
    list_axes[1].imshow(np.array(mask), cmap="gray")
    list_axes[0].set_title(titles[0], fontsize=15)
    list_axes[1].set_title(titles[1], fontsize=15)
    list_axes[2].imshow(np.array(result_img))
    list_axes[2].set_title(titles[2], fontsize=15)

    fig.subplots_adjust(wspace=0.01 if is_horizontal else 0.00, hspace=0.01 if is_horizontal else 0.1)
    fig.tight_layout()
    return fig


im_path = "./example_input.jpg"

# prepare input
model_input_size = [1024, 1024]
orig_im = np.array(Image.open(im_path))
orig_im_size = orig_im.shape[0:2]
image = preprocess_image(orig_im, model_input_size)

# inference
result = net(image)

# post process
result_image = postprocess_image(result[0][0], orig_im_size)

# save result
pil_im = Image.fromarray(result_image)
no_bg_image = Image.new("RGBA", pil_im.size, (0, 0, 0, 0))
orig_image = Image.open(im_path)
no_bg_image.paste(orig_image, mask=pil_im)
no_bg_image.save("example_image_no_bg.png")

visualize_result(orig_image, pil_im, no_bg_image);