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 import torch from torchvision import transforms import requests from io import BytesIO import zipfile # Fix the HF space permission error when using from_pretrained(..., trust_remote_code=True) os.environ["HF_MODULES_CACHE"] = os.path.join("/tmp/hf_cache", "modules") import transformers transformers.utils.move_cache() torch.set_float32_matmul_precision('high') torch.jit.script = lambda f: f device = "cuda" if torch.cuda.is_available() else "cpu" ## 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 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', 'Matting-dynamic': 'BiRefNet_dynamic-matting', } 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_simple(image): assert (image is not None), 'AssertionError: images cannot be None.' global birefnet # Use best default settings _weights_file = '/'.join(('zhengpeng7', usage_to_weights_file['General'])) birefnet = transformers.AutoModelForImageSegmentation.from_pretrained(_weights_file, trust_remote_code=True) birefnet.to(device) birefnet.eval(); birefnet.half() # Auto resolution - use 1024x1024 for best results resolution = (1024, 1024) image_ori = Image.fromarray(image) image = image_ori.convert('RGB') # Preprocess the image 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().cpu() pred = preds[0].squeeze() # Show Results pred_pil = transforms.ToPILImage()(pred) image_masked = refine_foreground(image, pred_pil, device=device) image_masked.putalpha(pred_pil.resize(image.size)) torch.cuda.empty_cache() return (image_masked, image_ori) examples = [[_] for _ in glob('examples/*')][:] # Custom CSS for styled buttons custom_css = """ .submit-btn { background: linear-gradient(45deg, #ff6b6b, #ee5a5a) !important; color: white !important; border: none !important; border-radius: 12px !important; padding: 12px 24px !important; font-weight: 600 !important; transition: all 0.3s ease !important; box-shadow: 0 4px 15px rgba(255, 107, 107, 0.3) !important; } .submit-btn:hover { background: linear-gradient(45deg, #ff5252, #d32f2f) !important; transform: translateY(-2px) !important; box-shadow: 0 6px 20px rgba(255, 107, 107, 0.4) !important; } .clear-btn { background: linear-gradient(45deg, #64b5f6, #42a5f5) !important; color: white !important; border: none !important; border-radius: 12px !important; padding: 12px 24px !important; font-weight: 600 !important; transition: all 0.3s ease !important; box-shadow: 0 4px 15px rgba(100, 181, 246, 0.3) !important; } .clear-btn:hover { background: linear-gradient(45deg, #2196f3, #1976d2) !important; transform: translateY(-2px) !important; box-shadow: 0 6px 20px rgba(100, 181, 246, 0.4) !important; } """ demo = gr.Interface( fn=predict_simple, inputs=[ gr.Image(label='Upload an image') ], outputs=gr.ImageSlider(label="Result", type="pil", format='png'), examples=examples, css=custom_css, submit_btn=gr.Button("🚀 Submit", elem_classes="submit-btn"), clear_btn=gr.Button("🗑️ Clear", elem_classes="clear-btn"), ) if __name__ == "__main__": demo.launch(debug=True)