Spaces:
Sleeping
Sleeping
File size: 2,177 Bytes
273eb2f |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 |
import gradio as gr
from PIL import Image
import torch
from torchvision import transforms
from transformers import AutoModelForImageSegmentation
# Setup constants
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
# Define image transformation pipeline
transform_image = transforms.Compose([
transforms.Resize((1024, 1024)),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])
# Load the model ONCE globally
try:
torch.set_float32_matmul_precision("high")
model = AutoModelForImageSegmentation.from_pretrained(
"ZhengPeng7/BiRefNet_lite",
trust_remote_code=True
).to(DEVICE)
print("Model loaded successfully.")
except Exception as e:
print(f"Error loading model: {str(e)}")
model = None
def process_image(image):
"""Process a single image and remove its background"""
image = image.convert("RGB")
original_size = image.size
input_tensor = transform_image(image).unsqueeze(0).to(DEVICE)
with torch.no_grad():
preds = model(input_tensor)[-1].sigmoid().cpu()
pred = preds[0].squeeze()
mask = transforms.ToPILImage()(pred).resize(original_size)
result = image.copy()
result.putalpha(mask)
return result
def predict(image):
"""Gradio interface function"""
if model is None:
raise gr.Error("Model not loaded. Check server logs.")
if image is None:
return None, None # Return None for both image and file
try:
result_image = process_image(image)
file_path = "processed_image.png"
result_image.save(file_path, "PNG")
return result_image, file_path
except Exception as e:
raise gr.Error(f"Error processing image: {e}")
# Gradio interface
interface = gr.Interface(
fn=predict,
inputs=gr.Image(type="pil"),
outputs=[
gr.Image(type="pil", label="Processed Image"),
gr.File(label="Download Processed Image")
],
examples=[['example.jpeg']],
title="Background Removal App",
description="Upload an image to remove its background and download the processed image as a PNG."
)
interface.launch()
|