import gradio as gr from loadimg import load_img import spaces from transformers import AutoModelForImageSegmentation import torch from torchvision import transforms from typing import Union, Tuple from PIL import Image, ImageDraw torch.set_float32_matmul_precision("high") DEVICE = "cuda" if torch.cuda.is_available() else "cpu" birefnet = AutoModelForImageSegmentation.from_pretrained( "ZhengPeng7/BiRefNet", trust_remote_code=True ) birefnet.to(DEVICE) birefnet.eval() # Keep inference inputs aligned with the loaded model precision (fp16/fp32). MODEL_DTYPE = next(birefnet.parameters()).dtype transform_image = transforms.Compose( [ transforms.Resize((1024, 1024)), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]), ] ) def fn(image: Union[Image.Image, str]) -> Tuple[Image.Image, Image.Image]: """ Remove the background from an image and return both the transparent version and the original. This function performs background removal using a BiRefNet segmentation model. It is intended for use with image input (either uploaded or from a URL). The function returns a transparent PNG version of the image with the background removed, along with the original RGB version for comparison. Args: image (PIL.Image or str): The input image, either as a PIL object or a filepath/URL string. Returns: tuple: - origin (PIL.Image): The original RGB image, unchanged. - processed_image (PIL.Image): The input image with the background removed and transparency applied. """ im = load_img(image, output_type="pil") im = im.convert("RGB") origin = im.copy() processed_image = process(im) return (origin, processed_image) @spaces.GPU(duration=20) def process(image: Image.Image) -> Image.Image: """ Apply BiRefNet-based image segmentation to remove the background. This function preprocesses the input image, runs it through a BiRefNet segmentation model to obtain a mask, and applies the mask as an alpha (transparency) channel to the original image. Args: image (PIL.Image): The input RGB image. Returns: PIL.Image: The image with the background removed, using the segmentation mask as transparency. """ image_size = image.size input_images = ( transform_image(image).unsqueeze(0).to(device=DEVICE, dtype=MODEL_DTYPE) ) # Prediction with torch.no_grad(): preds = birefnet(input_images)[-1].sigmoid().cpu() pred = preds[0].squeeze() pred_pil = transforms.ToPILImage()(pred) mask = pred_pil.resize(image_size) image.putalpha(mask) return image def process_file(f: str) -> str: """ Load an image file from disk, remove the background, and save the output as a transparent PNG. Args: f (str): Filepath of the image to process. Returns: str: Path to the saved PNG image with background removed. """ name_path = f.rsplit(".", 1)[0] + ".png" im = load_img(f, output_type="pil") im = im.convert("RGB") transparent = process(im) transparent.save(name_path) return name_path def _draw_bbox(image: Image.Image) -> Image.Image: """Draw a red bounding box around the non-transparent (foreground) region.""" alpha = image.split()[-1] bbox = alpha.getbbox() if bbox is None: return image boxed = image.copy() draw = ImageDraw.Draw(boxed) draw.rectangle(bbox, outline=(255, 0, 0, 255), width=3) return boxed def remove_bg_api( image: Image.Image, show_box: bool, canvas_size: float, padding_px: float, ) -> Image.Image: """ API endpoint used by the CRM Image Editor browser extension. Matches the payload the extension sends to POST /gradio_api/call/remove_bg: data: [image (FileData), show_box (bool), canvas_size (int), padding_px (int)] No resizing or padding is applied here — the returned image is the background-removed cutout at the original input resolution. `canvas_size` and `padding_px` are accepted (so the extension's existing payload shape still works) but intentionally ignored; do any resizing/padding on the extension side instead. Args: image: The uploaded product image. show_box: If True, draws a red bounding box around the detected foreground (useful for debugging). canvas_size: Unused — accepted for payload compatibility only. padding_px: Unused — accepted for payload compatibility only. Returns: PIL.Image: RGBA image with background removed, at original resolution. """ im = image.convert("RGB") cutout = process(im) if show_box: cutout = _draw_bbox(cutout) return cutout slider1 = gr.ImageSlider(label="Processed Image", type="pil", format="png") slider2 = gr.ImageSlider(label="Processed Image from URL", type="pil", format="png") image_upload = gr.Image(label="Upload an image") image_file_upload = gr.Image(label="Upload an image", type="filepath") url_input = gr.Textbox(label="Paste an image URL") output_file = gr.File(label="Output PNG File") # Example images chameleon = load_img("butterfly.jpg", output_type="pil") url_example = ( "https://hips.hearstapps.com/hmg-prod/images/gettyimages-1229892983-square.jpg" ) tab1 = gr.Interface( fn, inputs=image_upload, outputs=slider1, examples=[chameleon], api_name="image", description="""In case you are using an MCP, it is recommended you use the one from https://huggingface.co/spaces/hf-applications/background-removal or from the current app use the `png` api from the 'File Output' tab.""", ) tab2 = gr.Interface( fn, inputs=url_input, outputs=slider2, examples=[url_example], api_name="text", description="""In case you are using an MCP, it is recommended you use the one from https://huggingface.co/spaces/hf-applications/background-removal or from the current app use the `png` api from the 'File Output' tab.""", ) tab3 = gr.Interface( process_file, inputs=image_file_upload, outputs=output_file, examples=["butterfly.jpg"], api_name="png", description="""In case you are using an MCP, it is recommended you use the one from https://huggingface.co/spaces/hf-applications/background-removal or from the current app use the `png` api from the current tab""", ) tab4 = gr.Interface( remove_bg_api, inputs=[ gr.Image(label="Upload an image", type="pil"), gr.Checkbox(label="Show bounding box", value=False), gr.Number(label="Canvas size (px)", value=1200, precision=0), gr.Number(label="Padding (px)", value=0, precision=0), ], outputs=gr.Image(label="Result", type="pil", format="png"), api_name="remove_bg", description="Used by the CRM Image Editor browser extension. Returns a transparent PNG cutout at the original image resolution — no resizing or padding is applied server-side.", ) demo = gr.TabbedInterface( [tab1, tab2, tab3, tab4], ["Image Upload", "URL Input", "File Output", "Extension API"], title="Background Removal Tool", ) if __name__ == "__main__": demo.launch(show_error=True, mcp_server=True)