ZhengPeng7/BiRefNet_lite

1# These HF deployment codes refer to https://huggingface.co/not-lain/BiRefNet/raw/main/handler.py.2from typing import Dict, List, Any, Tuple3import os4import requests5from io import BytesIO6import cv27import numpy as np8from PIL import Image9import torch10from torchvision import transforms11from transformers import AutoModelForImageSegmentation12 13torch.set_float32_matmul_precision(["high", "highest"][0])14 15device = "cuda" if torch.cuda.is_available() else "cpu"16 17### image_proc.py18def refine_foreground(image, mask, r=90):19    if mask.size != image.size:20        mask = mask.resize(image.size)21    image = np.array(image) / 255.022    mask = np.array(mask) / 255.023    estimated_foreground = FB_blur_fusion_foreground_estimator_2(image, mask, r=r)24    image_masked = Image.fromarray((estimated_foreground * 255.0).astype(np.uint8))25    return image_masked26 27 28def FB_blur_fusion_foreground_estimator_2(image, alpha, r=90):29    # Thanks to the source: https://github.com/Photoroom/fast-foreground-estimation30    alpha = alpha[:, :, None]31    F, blur_B = FB_blur_fusion_foreground_estimator(image, image, image, alpha, r)32    return FB_blur_fusion_foreground_estimator(image, F, blur_B, alpha, r=6)[0]33 34 35def FB_blur_fusion_foreground_estimator(image, F, B, alpha, r=90):36    if isinstance(image, Image.Image):37        image = np.array(image) / 255.038    blurred_alpha = cv2.blur(alpha, (r, r))[:, :, None]39 40    blurred_FA = cv2.blur(F * alpha, (r, r))41    blurred_F = blurred_FA / (blurred_alpha + 1e-5)42 43    blurred_B1A = cv2.blur(B * (1 - alpha), (r, r))44    blurred_B = blurred_B1A / ((1 - blurred_alpha) + 1e-5)45    F = blurred_F + alpha * \46        (image - alpha * blurred_F - (1 - alpha) * blurred_B)47    F = np.clip(F, 0, 1)48    return F, blurred_B49 50 51class ImagePreprocessor():52    def __init__(self, resolution: Tuple[int, int] = (1024, 1024)) -> None:53        self.transform_image = transforms.Compose([54            transforms.Resize(resolution),55            transforms.ToTensor(),56            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),57        ])58 59    def proc(self, image: Image.Image) -> torch.Tensor:60        image = self.transform_image(image)61        return image62 63usage_to_weights_file = {64    'General': 'BiRefNet',65    'General-HR': 'BiRefNet_HR',66    'General-Lite': 'BiRefNet_lite',67    'General-Lite-2K': 'BiRefNet_lite-2K',68    'General-reso_512': 'BiRefNet-reso_512',69    'Matting': 'BiRefNet-matting',70    'Portrait': 'BiRefNet-portrait',71    'DIS': 'BiRefNet-DIS5K',72    'HRSOD': 'BiRefNet-HRSOD',73    'COD': 'BiRefNet-COD',74    'DIS-TR_TEs': 'BiRefNet-DIS5K-TR_TEs',75    'General-legacy': 'BiRefNet-legacy'76}77 78# Choose the version of BiRefNet here.79usage = 'General-Lite'80 81# Set resolution82if usage in ['General-Lite-2K']:83    resolution = (2560, 1440)84elif usage in ['General-reso_512']:85    resolution = (512, 512)86elif usage in ['General-HR']:87    resolution = (2048, 2048)88else:89    resolution = (1024, 1024) 90 91half_precision = True92 93class EndpointHandler():94    def __init__(self, path=''):95        self.birefnet = AutoModelForImageSegmentation.from_pretrained(96            '/'.join(('zhengpeng7', usage_to_weights_file[usage])), trust_remote_code=True97        )98        self.birefnet.to(device)99        self.birefnet.eval()100        if half_precision:101            self.birefnet.half()102 103    def __call__(self, data: Dict[str, Any]):104        """105        data args:106            inputs (:obj: `str`)107            date (:obj: `str`)108        Return:109            A :obj:`list` | `dict`: will be serialized and returned110        """111        print('data["inputs"] = ', data["inputs"])112        image_src = data["inputs"]113        if isinstance(image_src, str):114            if os.path.isfile(image_src):115                image_ori = Image.open(image_src)116            else:117                response = requests.get(image_src)118                image_data = BytesIO(response.content)119                image_ori = Image.open(image_data)120        else:121            image_ori = Image.fromarray(image_src)122 123        image = image_ori.convert('RGB')124        # Preprocess the image125        image_preprocessor = ImagePreprocessor(resolution=tuple(resolution))126        image_proc = image_preprocessor.proc(image)127        image_proc = image_proc.unsqueeze(0)128 129        # Prediction130        with torch.no_grad():131            preds = self.birefnet(image_proc.to(device).half() if half_precision else image_proc.to(device))[-1].sigmoid().cpu()132        pred = preds[0].squeeze()133 134        # Show Results135        pred_pil = transforms.ToPILImage()(pred)136        image_masked = refine_foreground(image, pred_pil)137        image_masked.putalpha(pred_pil.resize(image.size))138        return image_masked139