| import torch |
| from transformers import AutoModelForImageSegmentation |
| from torchvision import transforms |
| from PIL import Image |
| import numpy as np |
|
|
| |
| _device = torch.device("cuda" if torch.cuda.is_available() else "cpu") |
|
|
| |
| _birefnet = None |
| _birefnet_lite = None |
|
|
| |
| _transform = transforms.Compose( |
| [ |
| transforms.Resize((768, 768)), |
| transforms.ToTensor(), |
| transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]), |
| ] |
| ) |
|
|
|
|
| def load_model(fast_mode=False): |
| global _birefnet, _birefnet_lite |
|
|
| if fast_mode: |
| if _birefnet_lite is None: |
| _birefnet_lite = AutoModelForImageSegmentation.from_pretrained( |
| "ZhengPeng7/BiRefNet_lite", trust_remote_code=True |
| ).to(_device) |
| return _birefnet_lite |
|
|
| else: |
| if _birefnet is None: |
| _birefnet = AutoModelForImageSegmentation.from_pretrained( |
| "ZhengPeng7/BiRefNet", trust_remote_code=True |
| ).to(_device) |
| return _birefnet |
|
|
|
|
| def unload_model(): |
| global _birefnet, _birefnet_lite |
| _birefnet = None |
| _birefnet_lite = None |
| if torch.cuda.is_available(): |
| torch.cuda.empty_cache() |
|
|
|
|
| def _apply_mask(image: Image.Image, mask_tensor, bg, bg_type: str): |
| size = image.size |
| mask = transforms.ToPILImage()(mask_tensor.squeeze()).resize(size) |
|
|
| if bg_type == "Color" and isinstance(bg, str) and bg.startswith("#"): |
| color_rgb = tuple(int(bg[i : i + 2], 16) for i in (1, 3, 5)) |
| bg_image = Image.new("RGBA", size, color_rgb + (255,)) |
| elif bg_type in ["Image", "Video"]: |
| bg_image = bg.convert("RGBA").resize(size) |
| else: |
| raise ValueError("Invalid background type or format") |
|
|
| return Image.composite(image.convert("RGBA"), bg_image, mask) |
|
|
|
|
| def _get_transform(image_size: tuple[int, int] | None = None): |
| if image_size: |
| resize = transforms.Resize(image_size) |
| else: |
| resize = transforms.Resize((768, 768)) |
|
|
| return transforms.Compose( |
| [ |
| resize, |
| transforms.ToTensor(), |
| transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]), |
| ] |
| ) |
|
|
|
|
| def predict( |
| input: Image.Image | list[Image.Image], |
| bg: str | Image.Image, |
| bg_type: str = "Color", |
| fast_mode: bool = True, |
| bg_frames: list[np.ndarray] | None = None, |
| image_size_mode: str = "default", |
| custom_size: tuple[int, int] | None = None, |
| ) -> Image.Image | list[Image.Image]: |
|
|
| model = _birefnet_lite if fast_mode else _birefnet |
| if model is None: |
| raise RuntimeError("Model not loaded. Please call `load_model()` first.") |
|
|
| def get_transform_for_img(img: Image.Image): |
| if image_size_mode == "origin": |
| return _get_transform(img.size) |
| elif image_size_mode == "custom": |
| if not custom_size: |
| raise ValueError("`custom_size` must be provided for 'custom' mode.") |
| return _get_transform(custom_size) |
| else: |
| return _get_transform() |
|
|
| if isinstance(input, Image.Image): |
| _transform = get_transform_for_img(input) |
| input_tensor = _transform(input).unsqueeze(0).to(_device) |
| with torch.no_grad(): |
| pred = model(input_tensor)[-1].sigmoid().cpu() |
| bg_image = ( |
| Image.fromarray(bg_frames[0]) if bg_type == "Video" and bg_frames else bg |
| ) |
| return _apply_mask(input, pred[0], bg_image, bg_type) |
|
|
| elif isinstance(input, list): |
| results = [] |
| for idx, frame in enumerate(input): |
| _transform = get_transform_for_img(frame) |
| input_tensor = _transform(frame).unsqueeze(0).to(_device) |
| with torch.no_grad(): |
| pred = model(input_tensor)[-1].sigmoid().cpu() |
| if bg_type == "Video" and bg_frames: |
| bg_image = Image.fromarray(bg_frames[idx % len(bg_frames)]) |
| else: |
| bg_image = bg |
| results.append(_apply_mask(frame, pred[0], bg_image, bg_type)) |
| return results |
|
|
| else: |
| raise TypeError("`input` must be a PIL.Image or a list of PIL.Image.") |
|
|
|
|
| |
| |
|
|
| |
| |
|
|
| |
| |
|
|
