Spaces:
Running
Running
| import os | |
| import cv2 | |
| import numpy as np | |
| import torch | |
| import gradio as gr | |
| # import spaces | |
| from glob import glob | |
| from typing import Tuple | |
| from PIL import Image | |
| # from gradio_imageslider import ImageSlider | |
| import transformers | |
| import torch | |
| from torchvision import transforms | |
| import requests | |
| from io import BytesIO | |
| import zipfile | |
| from rankseg import RankSEG | |
| torch.set_float32_matmul_precision('high') | |
| # torch.jit.script = lambda f: f | |
| device = "cuda" if torch.cuda.is_available() else "cpu" | |
| RANKSEG_METRICS = ["dice", "iou"] | |
| def rgba2rgb(img): | |
| """ | |
| Convert RGBA image to RGB with white background. | |
| Supports both PIL.Image and numpy.ndarray. | |
| """ | |
| # 1. Handle PIL Image | |
| if isinstance(img, Image.Image): | |
| img = img.convert("RGBA") | |
| bg = Image.new("RGBA", img.size, (255, 255, 255)) | |
| return Image.alpha_composite(bg, img).convert("RGB") | |
| # 2. Handle Numpy Array (OpenCV) | |
| elif isinstance(img, np.ndarray): | |
| # Grayscale to RGB | |
| if img.ndim == 2: | |
| return np.stack([img] * 3, axis=-1) | |
| # Already 3 channels | |
| if img.shape[2] == 3: | |
| return img | |
| # RGBA to RGB (blending with white) | |
| elif img.shape[2] == 4: | |
| # Normalize alpha to 0-1 and keep shape (H, W, 1) | |
| alpha = img[..., 3:4].astype(float) / 255.0 | |
| foreground = img[..., :3].astype(float) | |
| background = 255.0 | |
| # Blend formula: source * alpha + bg * (1 - alpha) | |
| out = foreground * alpha + background * (1.0 - alpha) | |
| return out.clip(0, 255).astype(np.uint8) | |
| else: | |
| raise TypeError(f"Unsupported type: {type(img)}") | |
| ## CPU version refinement | |
| def FB_blur_fusion_foreground_estimator_cpu(image, FG, B, alpha, r=90): | |
| if isinstance(image, Image.Image): | |
| image = np.array(image) / 255.0 | |
| blurred_alpha = cv2.blur(alpha, (r, r))[:, :, None] | |
| blurred_FGA = cv2.blur(FG * alpha, (r, r)) | |
| blurred_FG = blurred_FGA / (blurred_alpha + 1e-5) | |
| blurred_B1A = cv2.blur(B * (1 - alpha), (r, r)) | |
| blurred_B = blurred_B1A / ((1 - blurred_alpha) + 1e-5) | |
| FG = blurred_FG + alpha * (image - alpha * blurred_FG - (1 - alpha) * blurred_B) | |
| FG = np.clip(FG, 0, 1) | |
| return FG, blurred_B | |
| def FB_blur_fusion_foreground_estimator_cpu_2(image, alpha, r=90): | |
| # Thanks to the source: https://github.com/Photoroom/fast-foreground-estimation | |
| alpha = alpha[:, :, None] | |
| FG, blur_B = FB_blur_fusion_foreground_estimator_cpu(image, image, image, alpha, r) | |
| return FB_blur_fusion_foreground_estimator_cpu(image, FG, blur_B, alpha, r=6)[0] | |
| ## GPU version refinement | |
| def mean_blur(x, kernel_size): | |
| """ | |
| equivalent to cv.blur | |
| x: [B, C, H, W] | |
| """ | |
| if kernel_size % 2 == 0: | |
| pad_l = kernel_size // 2 - 1 | |
| pad_r = kernel_size // 2 | |
| pad_t = kernel_size // 2 - 1 | |
| pad_b = kernel_size // 2 | |
| else: | |
| pad_l = pad_r = pad_t = pad_b = kernel_size // 2 | |
| x_padded = torch.nn.functional.pad(x, (pad_l, pad_r, pad_t, pad_b), mode='replicate') | |
| return torch.nn.functional.avg_pool2d(x_padded, kernel_size=(kernel_size, kernel_size), stride=1, count_include_pad=False) | |
| def FB_blur_fusion_foreground_estimator_gpu(image, FG, B, alpha, r=90): | |
| as_dtype = lambda x, dtype: x.to(dtype) if x.dtype != dtype else x | |
| input_dtype = image.dtype | |
| # convert image to float to avoid overflow | |
| image = as_dtype(image, torch.float32) | |
| FG = as_dtype(FG, torch.float32) | |
| B = as_dtype(B, torch.float32) | |
| alpha = as_dtype(alpha, torch.float32) | |
| blurred_alpha = mean_blur(alpha, kernel_size=r) | |
| blurred_FGA = mean_blur(FG * alpha, kernel_size=r) | |
| blurred_FG = blurred_FGA / (blurred_alpha + 1e-5) | |
| blurred_B1A = mean_blur(B * (1 - alpha), kernel_size=r) | |
| blurred_B = blurred_B1A / ((1 - blurred_alpha) + 1e-5) | |
| FG_output = blurred_FG + alpha * (image - alpha * blurred_FG - (1 - alpha) * blurred_B) | |
| FG_output = torch.clamp(FG_output, 0, 1) | |
| return as_dtype(FG_output, input_dtype), as_dtype(blurred_B, input_dtype) | |
| def FB_blur_fusion_foreground_estimator_gpu_2(image, alpha, r=90): | |
| # Thanks to the source: https://github.com/ZhengPeng7/BiRefNet/issues/226#issuecomment-3016433728 | |
| FG, blur_B = FB_blur_fusion_foreground_estimator_gpu(image, image, image, alpha, r) | |
| return FB_blur_fusion_foreground_estimator_gpu(image, FG, blur_B, alpha, r=6)[0] | |
| def refine_foreground(image, mask, r=90, device='cuda'): | |
| """both image and mask are in range of [0, 1]""" | |
| if mask.size != image.size: | |
| mask = mask.resize(image.size) | |
| if device == 'cuda': | |
| image = transforms.functional.to_tensor(image).float().cuda() | |
| mask = transforms.functional.to_tensor(mask).float().cuda() | |
| image = image.unsqueeze(0) | |
| mask = mask.unsqueeze(0) | |
| estimated_foreground = FB_blur_fusion_foreground_estimator_gpu_2(image, mask, r=r) | |
| estimated_foreground = estimated_foreground.squeeze() | |
| estimated_foreground = (estimated_foreground.mul(255.0)).to(torch.uint8) | |
| estimated_foreground = estimated_foreground.permute(1, 2, 0).contiguous().cpu().numpy().astype(np.uint8) | |
| else: | |
| image = np.array(image, dtype=np.float32) / 255.0 | |
| mask = np.array(mask, dtype=np.float32) / 255.0 | |
| estimated_foreground = FB_blur_fusion_foreground_estimator_cpu_2(image, mask, r=r) | |
| estimated_foreground = (estimated_foreground * 255.0).astype(np.uint8) | |
| estimated_foreground = Image.fromarray(np.ascontiguousarray(estimated_foreground)) | |
| return estimated_foreground | |
| def get_rankseg_pred(pred: torch.Tensor, metric: str) -> torch.Tensor: | |
| # BiRefNet produces a single foreground probability map, so RankSEG should | |
| # return a binary mask for that one channel instead of a multiclass map. | |
| rankseg = RankSEG(metric=metric, output_mode='multilabel', solver='RMA') | |
| probs = pred.unsqueeze(0).unsqueeze(0).to(torch.float32) | |
| return rankseg.predict(probs).squeeze(0).squeeze(0).to(torch.float32) | |
| def get_soft_gate(rankseg_mask: torch.Tensor, dilate_kernel: int = 9, blur_kernel: int = 15) -> torch.Tensor: | |
| support = rankseg_mask.unsqueeze(0).unsqueeze(0).to(torch.float32) | |
| dilated = torch.nn.functional.max_pool2d( | |
| support, | |
| kernel_size=dilate_kernel, | |
| stride=1, | |
| padding=dilate_kernel // 2, | |
| ) | |
| soft_gate = torch.nn.functional.avg_pool2d( | |
| dilated, | |
| kernel_size=blur_kernel, | |
| stride=1, | |
| padding=blur_kernel // 2, | |
| ) | |
| return soft_gate.squeeze(0).squeeze(0).clamp(0, 1) | |
| def build_hybrid_alpha(pred: torch.Tensor, rankseg_mask: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: | |
| soft_gate = get_soft_gate(rankseg_mask) | |
| hard_alpha = (pred * rankseg_mask.to(torch.float32)).clamp(0, 1) | |
| soft_alpha = torch.where(rankseg_mask > 0, pred, pred * soft_gate).clamp(0, 1) | |
| return hard_alpha, soft_alpha | |
| def build_alpha_cutout(image: Image.Image, mask: Image.Image) -> Image.Image: | |
| output = image.copy() | |
| output.putalpha(mask.resize(image.size)) | |
| return output | |
| def build_masked_image(image: Image.Image, mask: Image.Image) -> Image.Image: | |
| refined = refine_foreground(image, mask, device=device) | |
| refined.putalpha(mask.resize(image.size)) | |
| return refined | |
| def load_image(image_src): | |
| if isinstance(image_src, str): | |
| if os.path.isfile(image_src): | |
| image_ori = Image.open(image_src) | |
| else: | |
| response = requests.get(image_src) | |
| response.raise_for_status() | |
| image_data = BytesIO(response.content) | |
| image_ori = Image.open(image_data) | |
| else: | |
| image_ori = Image.fromarray(image_src) | |
| if image_ori.mode == 'RGBA': | |
| image_ori = rgba2rgb(image_ori) | |
| return image_ori.convert('RGB') | |
| class ImagePreprocessor(): | |
| def __init__(self, resolution: Tuple[int, int] = (1024, 1024)) -> None: | |
| # Input resolution is on WxH. | |
| self.transform_image = transforms.Compose([ | |
| transforms.Resize(resolution[::-1]), | |
| transforms.ToTensor(), | |
| transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]), | |
| ]) | |
| def proc(self, image: Image.Image) -> torch.Tensor: | |
| image = self.transform_image(image) | |
| return image | |
| usage_to_weights_file = { | |
| 'General': 'BiRefNet', | |
| 'General-HR': 'BiRefNet_HR', | |
| 'Matting-HR': 'BiRefNet_HR-matting', | |
| 'Matting': 'BiRefNet-matting', | |
| 'Portrait': 'BiRefNet-portrait', | |
| 'General-reso_512': 'BiRefNet_512x512', | |
| 'General-Lite': 'BiRefNet_lite', | |
| 'General-Lite-2K': 'BiRefNet_lite-2K', | |
| 'Anime-Lite': 'BiRefNet_lite-Anime', | |
| 'DIS': 'BiRefNet-DIS5K', | |
| 'HRSOD': 'BiRefNet-HRSOD', | |
| 'COD': 'BiRefNet-COD', | |
| 'DIS-TR_TEs': 'BiRefNet-DIS5K-TR_TEs', | |
| 'General-legacy': 'BiRefNet-legacy', | |
| 'General-dynamic': 'BiRefNet_dynamic', | |
| } | |
| birefnet = transformers.AutoModelForImageSegmentation.from_pretrained('/'.join(('zhengpeng7', usage_to_weights_file['General'])), trust_remote_code=True) | |
| birefnet.to(device) | |
| birefnet.eval(); birefnet.half() | |
| # @spaces.GPU | |
| def predict(images, resolution, weights_file, enable_rankseg, rankseg_metric): | |
| assert (images is not None), 'AssertionError: images cannot be None.' | |
| global birefnet | |
| # Load BiRefNet with chosen weights | |
| _weights_file = '/'.join(('zhengpeng7', usage_to_weights_file[weights_file] if weights_file is not None else usage_to_weights_file['General'])) | |
| print('Using weights: {}.'.format(_weights_file)) | |
| birefnet = transformers.AutoModelForImageSegmentation.from_pretrained(_weights_file, trust_remote_code=True) | |
| birefnet.to(device) | |
| birefnet.eval(); birefnet.half() | |
| try: | |
| resolution = [int(int(reso)//32*32) for reso in resolution.strip().split('x')] | |
| except: | |
| if weights_file in ['General-HR', 'Matting-HR']: | |
| resolution = (2048, 2048) | |
| elif weights_file in ['General-Lite-2K']: | |
| resolution = (2560, 1440) | |
| elif weights_file in ['General-reso_512']: | |
| resolution = (512, 512) | |
| else: | |
| if weights_file in ['General-dynamic']: | |
| resolution = None | |
| print('Using the original size (div by 32) for inference.') | |
| else: | |
| resolution = (1024, 1024) | |
| print('Invalid resolution input. Automatically changed to 1024x1024 / 2048x2048 / 2560x1440.') | |
| if isinstance(images, list): | |
| raw_save_paths = [] | |
| rankseg_hard_save_paths = [] | |
| rankseg_soft_save_paths = [] | |
| save_dir = 'preds-BiRefNet' | |
| if not os.path.exists(save_dir): | |
| os.makedirs(save_dir) | |
| tab_is_batch = True | |
| else: | |
| images = [images] | |
| tab_is_batch = False | |
| rankseg_metric = (rankseg_metric or 'dice').lower() | |
| if rankseg_metric not in RANKSEG_METRICS: | |
| rankseg_metric = 'dice' | |
| for image_src in images: | |
| image = load_image(image_src) | |
| # Preprocess the image | |
| if resolution is None: | |
| resolution_div_by_32 = [int(int(reso)//32*32) for reso in image.size] | |
| if resolution_div_by_32 != resolution: | |
| resolution = resolution_div_by_32 | |
| image_preprocessor = ImagePreprocessor(resolution=tuple(resolution)) | |
| image_proc = image_preprocessor.proc(image) | |
| image_proc = image_proc.unsqueeze(0) | |
| # Prediction | |
| with torch.no_grad(): | |
| preds = birefnet(image_proc.to(device).half())[-1].sigmoid().float().cpu() | |
| pred = preds[0].squeeze() | |
| pred_pil = transforms.ToPILImage()(pred) | |
| raw_image_masked = build_alpha_cutout(image, pred_pil) | |
| rankseg_hard_image_masked = None | |
| rankseg_soft_image_masked = None | |
| if enable_rankseg: | |
| rankseg_pred = get_rankseg_pred(pred, rankseg_metric) | |
| hard_alpha, soft_alpha = build_hybrid_alpha(pred, rankseg_pred) | |
| rankseg_hard_image_masked = build_alpha_cutout(image, transforms.ToPILImage()(hard_alpha)) | |
| rankseg_soft_image_masked = build_alpha_cutout(image, transforms.ToPILImage()(soft_alpha)) | |
| if device == 'cuda': | |
| torch.cuda.empty_cache() | |
| if tab_is_batch: | |
| image_name = os.path.splitext(os.path.basename(image_src))[0] | |
| raw_save_file_path = os.path.join(save_dir, f"{image_name}_pred.png") | |
| raw_image_masked.save(raw_save_file_path) | |
| raw_save_paths.append(raw_save_file_path) | |
| if enable_rankseg and rankseg_hard_image_masked is not None: | |
| rankseg_hard_save_file_path = os.path.join(save_dir, f"{image_name}_pred_rankseg.png") | |
| rankseg_hard_image_masked.save(rankseg_hard_save_file_path) | |
| rankseg_hard_save_paths.append(rankseg_hard_save_file_path) | |
| if enable_rankseg and rankseg_soft_image_masked is not None: | |
| rankseg_soft_save_file_path = os.path.join(save_dir, f"{image_name}_pred_softgate.png") | |
| rankseg_soft_image_masked.save(rankseg_soft_save_file_path) | |
| rankseg_soft_save_paths.append(rankseg_soft_save_file_path) | |
| if tab_is_batch: | |
| zip_file_path = os.path.join(save_dir, "{}.zip".format(save_dir)) | |
| with zipfile.ZipFile(zip_file_path, 'w') as zipf: | |
| for file in raw_save_paths + rankseg_hard_save_paths + rankseg_soft_save_paths: | |
| zipf.write(file, os.path.basename(file)) | |
| return raw_save_paths, rankseg_hard_save_paths, rankseg_soft_save_paths, zip_file_path | |
| else: | |
| return image, raw_image_masked, rankseg_hard_image_masked, rankseg_soft_image_masked | |
| examples = [[_] for _ in glob('examples/*')][:] | |
| # Add the option of resolution in a text box. | |
| for idx_example, example in enumerate(examples): | |
| if 'My_' in example[0]: | |
| example_resolution = '2048x2048' | |
| model_choice = 'Matting-HR' | |
| else: | |
| example_resolution = '1024x1024' | |
| model_choice = 'General' | |
| examples[idx_example] = examples[idx_example] + [example_resolution, model_choice, True, 'dice'] | |
| examples_url = [ | |
| ['https://hips.hearstapps.com/hmg-prod/images/gettyimages-1229892983-square.jpg'], | |
| ] | |
| for idx_example_url, example_url in enumerate(examples_url): | |
| examples_url[idx_example_url] = examples_url[idx_example_url] + ['1024x1024', 'General', True, 'dice'] | |
| descriptions = ('Upload a picture, our model will extract a highly accurate segmentation of the subject in it.\n)' | |
| ' The resolution used in our training was `1024x1024`, which is the suggested resolution to obtain good results! `2048x2048` is suggested for BiRefNet_HR.\n' | |
| ' Our codes can be found at https://github.com/ZhengPeng7/BiRefNet.\n' | |
| ' We also maintain the HF model of BiRefNet at https://huggingface.co/ZhengPeng7/BiRefNet for easier access.') | |
| tab_image = gr.Interface( | |
| fn=predict, | |
| inputs=[ | |
| gr.Image(label='Upload an image', image_mode='RGBA'), # Keep alpha channel | |
| gr.Textbox(lines=1, placeholder="Type the resolution (`WxH`) you want, e.g., `1024x1024`.", label="Resolution"), | |
| gr.Radio(list(usage_to_weights_file.keys()), value='General', label="Weights", info="Choose the weights you want."), | |
| gr.Checkbox(value=True, label="Enable RankSEG"), | |
| gr.Radio(RANKSEG_METRICS, value='dice', label="RankSEG metric", info="Choose the target metric for RankSEG post-processing.") | |
| ], | |
| outputs=[ | |
| gr.Image(label="Original image", type="pil", format='png'), | |
| gr.Image(label="BiRefNet pred alpha", type="pil", format='png'), | |
| gr.Image(label="BiRefNet pred x RankSEG", type="pil", format='png'), | |
| gr.Image(label="BiRefNet soft-gated hybrid", type="pil", format='png'), | |
| ], | |
| examples=examples, | |
| api_name="image", | |
| description=descriptions, | |
| ) | |
| tab_text = gr.Interface( | |
| fn=predict, | |
| inputs=[ | |
| gr.Textbox(label="Paste an image URL"), | |
| gr.Textbox(lines=1, placeholder="Type the resolution (`WxH`) you want, e.g., `1024x1024`.", label="Resolution"), | |
| gr.Radio(list(usage_to_weights_file.keys()), value='General', label="Weights", info="Choose the weights you want."), | |
| gr.Checkbox(value=True, label="Enable RankSEG"), | |
| gr.Radio(RANKSEG_METRICS, value='dice', label="RankSEG metric", info="Choose the target metric for RankSEG post-processing.") | |
| ], | |
| outputs=[ | |
| gr.Image(label="Original image", type="pil", format='png'), | |
| gr.Image(label="BiRefNet pred alpha", type="pil", format='png'), | |
| gr.Image(label="BiRefNet pred x RankSEG", type="pil", format='png'), | |
| gr.Image(label="BiRefNet soft-gated hybrid", type="pil", format='png'), | |
| ], | |
| examples=examples_url, | |
| api_name="URL", | |
| description=descriptions+'\nTab-URL is partially modified from https://huggingface.co/spaces/not-lain/background-removal, thanks to this great work!', | |
| ) | |
| tab_batch = gr.Interface( | |
| fn=predict, | |
| inputs=[ | |
| gr.File(label="Upload multiple images", type="filepath", file_count="multiple"), | |
| gr.Textbox(lines=1, placeholder="Type the resolution (`WxH`) you want, e.g., `1024x1024`.", label="Resolution"), | |
| gr.Radio(list(usage_to_weights_file.keys()), value='General', label="Weights", info="Choose the weights you want."), | |
| gr.Checkbox(value=True, label="Enable RankSEG"), | |
| gr.Radio(RANKSEG_METRICS, value='dice', label="RankSEG metric", info="Choose the target metric for RankSEG post-processing.") | |
| ], | |
| outputs=[ | |
| gr.Gallery(label="BiRefNet pred alpha results"), | |
| gr.Gallery(label="BiRefNet pred x RankSEG results"), | |
| gr.Gallery(label="BiRefNet soft-gated hybrid results"), | |
| gr.File(label="Download masked images."), | |
| ], | |
| api_name="batch", | |
| description=descriptions+'\nTab-batch is partially modified from https://huggingface.co/spaces/NegiTurkey/Multi_Birefnetfor_Background_Removal, thanks to this great work!', | |
| ) | |
| demo = gr.TabbedInterface( | |
| [tab_image, tab_text, tab_batch], | |
| ['image', 'URL', 'batch'], | |
| title="Official Online Demo of BiRefNet", | |
| ) | |
| if __name__ == "__main__": | |
| demo.launch(debug=True) | |