Update app.py

app.py

@@ -1,64 +1,65 @@
+import os
 import gradio as gr
 import torch
 from torchvision import transforms
 from PIL import Image
 import requests
-from io import BytesIO
+from BiRefNet import BiRefNet
 
-# Download BiRefNet weights (automatically, if not present)
-model_url = "https://huggingface.co/ZhengPeng7/BiRefNet/resolve/main/BiRefNet.pth"
-model_path = "BiRefNet.pth"
+# 1. Download BiRefNet weights if not present
+MODEL_URL = "https://huggingface.co/ZhengPeng7/BiRefNet/resolve/main/BiRefNet.pth"
+MODEL_PATH = "BiRefNet.pth"
 
 def download_weights():
-    if not os.path.exists(model_path):
-        response = requests.get(model_url)
-        with open(model_path, "wb") as f:
-            f.write(response.content)
-
-# Define BiRefNet model loading (use standard repo's class definition)
-import os
+    if not os.path.exists(MODEL_PATH):
+        print("Downloading BiRefNet weights...")
+        r = requests.get(MODEL_URL)
+        with open(MODEL_PATH, "wb") as f:
+            f.write(r.content)
+        print("Done downloading BiRefNet weights.")
 
-class BiRefNetDummy(torch.nn.Module):
-    def forward(self, x):
-        # Dummy output, placeholder for real BiRefNet logic.
-        # Will be replaced by actual model code in the next step.
-        return torch.ones_like(x)
-
-ddef load_model():
+# 2. Load BiRefNet model
+def load_model():
     download_weights()
-    net = BiRefNet()
-    net.load_state_dict(torch.load("BiRefNet.pth", map_location="cpu"))
-    net.eval()
-    return net
+    model = BiRefNet()
+    state_dict = torch.load(MODEL_PATH, map_location="cpu")
+    model.load_state_dict(state_dict)
+    model.eval()
+    return model
 
 bi_ref_net = load_model()
 
+# 3. Define transforms (assuming model expects 224x224 or similar, adjust if needed)
+preprocess = transforms.Compose([
+    transforms.Resize((224, 224)),  # Adjust to BiRefNet input size if different
+    transforms.ToTensor()
+])
+
 def remove_bg(input_image):
-    # PIL image to tensor
-    transform = transforms.Compose([
-        transforms.ToTensor(),
-    ])
-    img_tensor = transform(input_image).unsqueeze(0)
+    # Preprocess image
+    image = input_image.convert("RGB")
+    img_tensor = preprocess(image).unsqueeze(0)  # Add batch dimension
+
+    # Inference (no gradients needed)
+    with torch.no_grad():
+        mask = bi_ref_net(img_tensor)[0, 0]   # Output mask from model, shape: [H, W]
 
-    # Get mask from BiRefNet (replace with actual model code later)
-    mask = bi_ref_net(img_tensor)
-    mask = mask.squeeze().detach().numpy()
-    mask_img = Image.fromarray((mask * 255).astype("uint8")).convert("L")
-    mask_img = mask_img.resize(input_image.size)
+    # Resize mask to original image size, normalize (if needed)
+    mask_img = transforms.ToPILImage()(mask.cpu().clamp(0, 1))
+    mask_img = mask_img.resize(image.size, Image.BILINEAR)
 
-    # Apply mask: set alpha channel
-    result = input_image.convert("RGBA")
+    # Create RGBA output by setting alpha to mask
+    result = image.convert("RGBA")
     result.putalpha(mask_img)
     return result
 
 demo = gr.Interface(
     fn=remove_bg,
     inputs=gr.Image(type="pil", label="Input Image"),
-    outputs=gr.Image(type="pil", label="Output (Background Removed)"),
+    outputs=gr.Image(type="pil", label="Background Removed (PNG)"),
     title="Backdrop Studio - BiRefNet Background Removal",
-    description="Upload an image to remove the background using BiRefNet AI. [v1.0 demo]"
+    description="Upload an image to remove the background using BiRefNet AI."
 )
 
 if __name__ == "__main__":
-    download_weights()  # Download weights if needed
-    demo.launch()
+    demo.launch()