import os import cv2 import numpy as np import torch import gradio as gr import spaces from glob import glob from typing import Optional, Tuple from PIL import Image from gradio_imageslider import ImageSlider from transformers import AutoModelForImageSegmentation from torchvision import transforms import requests from io import BytesIO import zipfile torch.set_float32_matmul_precision('high') torch.jit.script = lambda f: f device = "cuda" if torch.cuda.is_available() else "CPU" def array_to_pil_image(image: np.ndarray, size: Tuple[int, int] = (1024, 1024)) -> Image.Image: image = cv2.resize(image, size, interpolation=cv2.INTER_LINEAR) image = Image.fromarray(image).convert('RGB') return image class ImagePreprocessor(): def __init__(self, resolution: Tuple[int, int] = (1024, 1024)) -> None: self.transform_image = transforms.Compose([ # transforms.Resize(resolution), # 1. keep consistent with the cv2.resize used in training 2. redundant with that in path_to_image() 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-Lite': 'BiRefNet_lite', 'Portrait': 'BiRefNet-portrait', 'DIS': 'BiRefNet-DIS5K', 'HRSOD': 'BiRefNet-HRSOD', 'COD': 'BiRefNet-COD', 'DIS-TR_TEs': 'BiRefNet-DIS5K-TR_TEs', 'General-legacy': 'BiRefNet-legacy' } birefnet = AutoModelForImageSegmentation.from_pretrained('/'.join(('zhengpeng7', usage_to_weights_file['General'])), trust_remote_code=True) birefnet.to(device) birefnet.eval() @spaces.GPU def predict(images, resolution, weights_file): 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 = AutoModelForImageSegmentation.from_pretrained(_weights_file, trust_remote_code=True) birefnet.to(device) birefnet.eval() try: resolution = [int(int(reso)//32*32) for reso in resolution.strip().split('x')] except: resolution = [1024, 1024] print('Invalid resolution input. Automatically changed to 1024x1024.') if isinstance(images, list): # For tab_batch 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 for idx_image, image_src in enumerate(images): if isinstance(image_src, str): if os.path.isfile(image_src): image = np.array(Image.open(image_src)) else: response = requests.get(image_src) image_data = BytesIO(response.content) image = np.array(Image.open(image_data)) else: image = image_src image_shape = image.shape[:2] image_pil = array_to_pil_image(image, tuple(resolution)) # Preprocess the image image_preprocessor = ImagePreprocessor(resolution=tuple(resolution)) image_proc = image_preprocessor.proc(image_pil) image_proc = image_proc.unsqueeze(0) # Perform the prediction with torch.no_grad(): scaled_pred_tensor = birefnet(image_proc.to(device))[-1].sigmoid() if device == 'cuda': scaled_pred_tensor = scaled_pred_tensor.cpu() # Resize the prediction to match the original image shape pred = torch.nn.functional.interpolate(scaled_pred_tensor, size=image_shape, mode='bilinear', align_corners=True).squeeze().numpy() # Apply the prediction mask to the original image image_pil = image_pil.resize(pred.shape[::-1]) pred = np.repeat(np.expand_dims(pred, axis=-1), 3, axis=-1) image_masked = (pred * np.array(image_pil)).astype(np.uint8) torch.cuda.empty_cache() if tab_is_batch: save_file_path = os.path.join(save_dir, "{}.png".format(os.path.splitext(os.path.basename(image_src))[0])) cv2.imwrite(save_file_path, image_masked) save_paths.append(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 save_paths: zipf.write(file, os.path.basename(file)) return image, image_masked examples = [[_] for _ in glob('examples/*')][:] # Add the option of resolution in a text box. for idx_example, example in enumerate(examples): examples[idx_example].append('1024x1024') examples.append(examples[-1].copy()) examples[-1][1] = '512x512' 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].append('1024x1024') 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`, thus the suggested resolution to obtain good results!\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'), gr.Textbox(lines=1, placeholder="Type the resolution (`WxH`) you want, e.g., `1024x1024`. Higher resolutions can be much slower for inference.", label="Resolution"), gr.Radio(list(usage_to_weights_file.keys()), value='General', label="Weights", info="Choose the weights you want.") ], outputs=ImageSlider(label="BiRefNet's prediction", type="pil"), 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`. Higher resolutions can be much slower for inference.", label="Resolution"), gr.Radio(list(usage_to_weights_file.keys()), value='General', label="Weights", info="Choose the weights you want.") ], outputs=ImageSlider(label="BiRefNet's prediction", type="pil"), examples=examples_url, api_name="text", 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"), outputs=[gr.Gallery(label="BiRefNet's predictions"), 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', 'text', 'batch'], title="BiRefNet demo for subject extraction (general / salient / camouflaged / portrait).", ) if __name__ == "__main__": demo.launch(debug=True)