Spaces:
Sleeping
Sleeping
File size: 3,150 Bytes
c5b6b5f 5bdaefc c5b6b5f b814f78 5bdaefc b814f78 5bdaefc b814f78 c5b6b5f 5bdaefc c5b6b5f 5bdaefc c5b6b5f 5bdaefc c5b6b5f 5bdaefc c5b6b5f 5bdaefc c5b6b5f 684c85b 5eb497f bf469c8 c5b6b5f 5bdaefc 5eb497f |
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 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 |
"""
Background Removal API using rembg
Optimized for Hugging Face Spaces with GPU support
Simple Gradio-only implementation
"""
import gradio as gr
from rembg import remove, new_session
from PIL import Image
import io
import time
def create_rembg_session():
"""Create rembg session with GPU fallback to CPU if CUDA libraries are missing."""
print("🔥 Loading rembg model...")
start_time = time.time()
try:
session = new_session(model_name="u2net")
load_time = time.time() - start_time
print(f"✅ rembg model loaded in {load_time:.2f}s (GPU)")
return session
except Exception as e:
print(f"⚠️ GPU session failed, falling back to CPU: {e}")
try:
session = new_session(model_name="u2net", providers=["CPUExecutionProvider"])
load_time = time.time() - start_time
print(f"✅ rembg model loaded in {load_time:.2f}s (CPU)")
return session
except Exception as e2:
print(f"❌ Failed to load rembg: {e2}")
raise
# Global session for model reuse
rembg_session = create_rembg_session()
def remove_background(image):
"""
Remove background from image using rembg
Args:
image: PIL Image
Returns:
PIL Image with transparent background
"""
if image is None:
return None
print(f"📥 Received image: {image.size}, mode: {image.mode}")
try:
# Convert PIL Image to bytes
img_byte_arr = io.BytesIO()
image.save(img_byte_arr, format='PNG')
img_byte_arr.seek(0)
image_bytes = img_byte_arr.getvalue()
# Remove background using rembg with session reuse
print("🔄 Processing with rembg...")
start_time = time.time()
output_bytes = remove(image_bytes, session=rembg_session)
process_time = time.time() - start_time
print(f"✅ Background removed in {process_time:.2f}s")
# Convert back to PIL Image
output_image = Image.open(io.BytesIO(output_bytes)).convert("RGBA")
print(f"📤 Output image: {output_image.size}, mode: {output_image.mode}")
return output_image
except Exception as e:
print(f"❌ Error: {e}")
raise gr.Error(f"Background removal failed: {str(e)}")
# Create Gradio Interface
demo = gr.Interface(
fn=remove_background,
inputs=gr.Image(type="pil", label="Upload Image"),
outputs=gr.Image(type="pil", label="Background Removed"),
title="🎨 Background Removal API",
description="""
Fast background removal using rembg (u2net model).
**Features:**
- GPU-accelerated processing (when available)
- Session reuse for faster processing
- API access via gradio_client
**Python API Usage:**
```python
from gradio_client import Client
client = Client("https://tigger13-background-removal.hf.space")
result = client.predict("/path/to/image.png")
```
""",
examples=[]
)
if __name__ == "__main__":
demo.launch(
server_name="0.0.0.0",
server_port=7860,
show_error=True # Enable error reporting
)
|