Spaces:
Running on Zero
Running on Zero
| 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) | |
| 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) |