|
import io |
|
from pathlib import Path |
|
from typing import Union |
|
import numpy as np |
|
import cv2 |
|
|
|
import torch |
|
import torch.nn.functional as F |
|
from PIL import Image |
|
from rembg import remove |
|
from internals.data.task import ModelType |
|
|
|
import internals.util.image as ImageUtil |
|
from carvekit.api.high import HiInterface |
|
from internals.util.commons import download_image, read_url |
|
import onnxruntime as rt |
|
import huggingface_hub |
|
|
|
|
|
class RemoveBackground: |
|
def remove(self, image: Union[str, Image.Image]) -> Image.Image: |
|
if type(image) is str: |
|
image = Image.open(io.BytesIO(read_url(image))) |
|
|
|
output = remove(image) |
|
return output |
|
|
|
|
|
class RemoveBackgroundV2: |
|
def __init__(self): |
|
model_path = huggingface_hub.hf_hub_download("skytnt/anime-seg", "isnetis.onnx") |
|
self.anime_rembg = rt.InferenceSession( |
|
model_path, providers=["CUDAExecutionProvider", "CPUExecutionProvider"] |
|
) |
|
|
|
self.interface = HiInterface( |
|
object_type="object", |
|
batch_size_seg=5, |
|
batch_size_matting=1, |
|
device="cuda" if torch.cuda.is_available() else "cpu", |
|
seg_mask_size=640, |
|
matting_mask_size=2048, |
|
trimap_prob_threshold=231, |
|
trimap_dilation=30, |
|
trimap_erosion_iters=5, |
|
fp16=False, |
|
) |
|
|
|
def remove( |
|
self, image: Union[str, Image.Image], model_type: ModelType = ModelType.REAL |
|
) -> Image.Image: |
|
if type(image) is str: |
|
image = download_image(image) |
|
|
|
if model_type == ModelType.ANIME or model_type == ModelType.COMIC: |
|
print("Using Anime Background remover") |
|
_, img = self.__rmbg_fn(np.array(image)) |
|
|
|
return Image.fromarray(img) |
|
else: |
|
print("Using Real Background remover") |
|
img_path = Path.home() / ".cache" / "rm_bg.png" |
|
|
|
w, h = image.size |
|
if max(w, h) > 1536: |
|
image = ImageUtil.resize_image(image, dimension=1024) |
|
|
|
image.save(img_path) |
|
images_without_background = self.interface([img_path]) |
|
out = images_without_background[0] |
|
return out |
|
|
|
def __get_mask(self, img, s=1024): |
|
img = (img / 255).astype(np.float32) |
|
h, w = h0, w0 = img.shape[:-1] |
|
h, w = (s, int(s * w / h)) if h > w else (int(s * h / w), s) |
|
ph, pw = s - h, s - w |
|
img_input = np.zeros([s, s, 3], dtype=np.float32) |
|
img_input[ph // 2 : ph // 2 + h, pw // 2 : pw // 2 + w] = cv2.resize( |
|
img, (w, h) |
|
) |
|
img_input = np.transpose(img_input, (2, 0, 1)) |
|
img_input = img_input[np.newaxis, :] |
|
mask = self.anime_rembg.run(None, {"img": img_input})[0][0] |
|
mask = np.transpose(mask, (1, 2, 0)) |
|
mask = mask[ph // 2 : ph // 2 + h, pw // 2 : pw // 2 + w] |
|
mask = cv2.resize(mask, (w0, h0))[:, :, np.newaxis] |
|
return mask |
|
|
|
def __rmbg_fn(self, img): |
|
mask = self.__get_mask(img) |
|
img = (mask * img + 255 * (1 - mask)).astype(np.uint8) |
|
mask = (mask * 255).astype(np.uint8) |
|
img = np.concatenate([img, mask], axis=2, dtype=np.uint8) |
|
mask = mask.repeat(3, axis=2) |
|
return mask, img |
|
|
