Upload 4 files

for_gradio.py

@@ -0,0 +1,74 @@
+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()

requirements.txt

@@ -0,0 +1,9 @@
+# loadimg
+torch
+torchvision
+transformers
+kornia
+einops
+timm
+streamlit
+gradio

runtime.txt

@@ -0,0 +1 @@
+python_version=3.10