Spaces:

Janeka
/

pipeline2

Sleeping

App Files Files Community

Janeka commited on Apr 23, 2025

Commit

d554758

verified ·

1 Parent(s): 729a829

Update app.py

Browse files

Files changed (1) hide show

app.py +40 -93

app.py CHANGED Viewed

@@ -1,100 +1,47 @@
-import cv2
 import numpy as np
 import torch
 from PIL import Image
-from fastapi import FastAPI, UploadFile, File, HTTPException
-from fastapi.responses import HTMLResponse, Response
-import time
-from io import BytesIO
-import os
-from huggingface_hub import hf_hub_download
-# Manual implementation of BiRefNet components
-class BiRefNet(torch.nn.Module):
-    def __init__(self):
-        super().__init__()
-        # Simplified architecture for CPU
-        self.conv1 = torch.nn.Conv2d(4, 64, kernel_size=3, padding=1)
-        self.conv2 = torch.nn.Conv2d(64, 1, kernel_size=3, padding=1)
-    def forward(self, x):
-        x = torch.relu(self.conv1(x))
-        return torch.sigmoid(self.conv2(x))
-def inference_image(model, image, device="cpu"):
-    """Simplified inference for CPU"""
-    input_tensor = torch.from_numpy(image).permute(2,0,1).unsqueeze(0).float()/255.0
-    with torch.no_grad():
-        output = model(input_tensor.to(device))
-    return output.squeeze().cpu().numpy()
-app = FastAPI(title="BiRefNet Background Remover")
-# Configuration
-MAX_SIZE = 1024
-MODEL_PATH = "birefnet.pth"
-DEVICE = "cpu"
-# Initialize model
 model = BiRefNet()
-if not os.path.exists(MODEL_PATH):
-    hf_hub_download(repo_id="ZhengPeng7/BiRefNet",
-                   filename="birefnet.pth",
-                   local_dir=".",
-                   force_filename="birefnet.pth")
-model.load_state_dict(torch.load(MODEL_PATH, map_location=DEVICE))
-model.to(DEVICE)
 model.eval()
-def preprocess_image(image):
-    if image.shape[2] == 3:
-        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGBA)
-    h, w = image.shape[:2]
-    if max(h, w) > MAX_SIZE:
-        ratio = MAX_SIZE / max(h, w)
-        image = cv2.resize(image, (int(w*ratio), int(h*ratio)),
-                         interpolation=cv2.INTER_AREA)
-    return image
-@app.post("/remove_bg")
-async def remove_bg(file: UploadFile = File(...)):
-    try:
-        start_time = time.time()
-        contents = await file.read()
-        img = cv2.imdecode(np.frombuffer(contents, np.uint8), cv2.IMREAD_UNCHANGED)
-        img = preprocess_image(img)
-        with torch.no_grad():
-            alpha = inference_image(model, img, DEVICE)
-        if img.shape[2] == 3:
-            img = cv2.cvtColor(img, cv2.COLOR_RGB2RGBA)
-        img[:, :, 3] = (alpha * 255).astype(np.uint8)
-        img_bytes = BytesIO()
-        Image.fromarray(img).save(img_bytes, format="PNG")
-        return Response(
-            content=img_bytes.getvalue(),
-            media_type="image/png",
-            headers={"X-Processing-Time": f"{time.time()-start_time:.2f}s"}
-        )
-    except Exception as e:
-        raise HTTPException(500, f"Error: {str(e)}")
-@app.get("/", response_class=HTMLResponse)
-async def home():
-    return HTMLResponse("""
-    <html>
-    <body>
-        <h1>BiRefNet Background Remover</h1>
-        <form action="/remove_bg" method="post" enctype="multipart/form-data">
-            <input type="file" name="file" accept="image/*" required>
-            <button>Remove Background</button>
-        </form>
-    </body>
-    </html>
-    """)
-if __name__ == "__main__":
-    import uvicorn
-    uvicorn.run(app, host="0.0.0.0", port=7860)

+import gradio as gr
 import numpy as np
 import torch
+import cv2
 from PIL import Image
+from BiRefNet.models.BiRefNet import BiRefNet
+from BiRefNet.utils.dataloader import test_dataset
+# Initialize model (smaller version for CPU)
 model = BiRefNet()
+model.load_state_dict(torch.load('BiRefNet.pth', map_location=torch.device('cpu')))
 model.eval()
+def remove_background(input_image):
+    # Preprocess image
+    image = np.array(input_image)
+    image = cv2.resize(image, (320, 320))  # Smaller size for CPU
+    image = test_dataset.preprocess(image)
+    image = torch.from_numpy(image).unsqueeze(0)
+    # Inference
+    with torch.no_grad():
+        pred = model(image)
+    # Post-process
+    pred = pred.squeeze().cpu().numpy()
+    mask = (pred > 0.5).astype(np.uint8) * 255
+    # Apply mask to original image
+    original = cv2.resize(np.array(input_image), (320, 320))
+    result = cv2.bitwise_and(original, original, mask=mask)
+    return Image.fromarray(result), Image.fromarray(mask)
+# Gradio interface
+interface = gr.Interface(
+    fn=remove_background,
+    inputs=gr.Image(type="pil", label="Input Image"),
+    outputs=[
+        gr.Image(type="pil", label="Result"),
+        gr.Image(type="pil", label="Mask")
+    ],
+    title="BiRefNet Background Remover (CPU)",
+    description="Upload an image to remove the background. Works on CPU but may be slow."
+)
+interface.launch()