Spaces:
Build error
Build error
File size: 1,819 Bytes
e546fea 9bb32c5 808ad1a 958511f e546fea 808ad1a |
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 |
import torch
from PIL import Image
import numpy as np
import base64
import io
import gradio as gr
from your_model_imports import BiRefNet # replace with your actual model import
# Force CPU
device = torch.device("cpu")
# Load model
birefnet = BiRefNet() # or your model class
birefnet.to(device)
birefnet.eval() # set evaluation mode
# Helper to convert base64 to PIL
def b64_to_pil(b64_image):
header, data = b64_image.split(",", 1)
img_bytes = base64.b64decode(data)
return Image.open(io.BytesIO(img_bytes)).convert("RGBA")
# Helper to convert PIL to base64
def pil_to_b64(pil_img):
buffered = io.BytesIO()
pil_img.save(buffered, format="PNG")
img_str = base64.b64encode(buffered.getvalue()).decode("utf-8")
return f"data:image/png;base64,{img_str}"
# Background removal function
def remove_bg(image_b64):
try:
# Convert to PIL
img = b64_to_pil(image_b64)
# Convert PIL to tensor
img_tensor = torch.from_numpy(np.array(img)).permute(2,0,1).unsqueeze(0).float() / 255.0
img_tensor = img_tensor.to(device)
# Run model
with torch.no_grad():
output_tensor = birefnet(img_tensor)
# Convert output tensor to PIL
output_np = (output_tensor.squeeze().permute(1,2,0).numpy() * 255).astype(np.uint8)
output_pil = Image.fromarray(output_np)
# Convert to base64
return pil_to_b64(output_pil)
except Exception as e:
return f"ERROR: {str(e)}"
# Gradio interface
iface = gr.Interface(
fn=remove_bg,
inputs=gr.Image(type="pil", label="Input Image"),
outputs=gr.Image(type="auto", label="Background Removed"),
title="Background Remover Pixels",
description="Removes background using CPU-only model."
)
if __name__ == "__main__":
iface.launch()
|