Update app.py

app.py

@@ -1,35 +1,99 @@
 import os
-import subprocess
+import sys
+import requests
 import gradio as gr
 import numpy as np
 import torch
 import cv2
 from PIL import Image
-from huggingface_hub import hf_hub_download
+from tqdm import tqdm
 
-# Clone and install BiRefNet if not present
+# Download BiRefNet repository
 if not os.path.exists('BiRefNet'):
-    subprocess.run(['git', 'clone', 'https://github.com/ZhengPeng7/BiRefNet.git'])
-    os.chdir('BiRefNet')
-    subprocess.run(['pip', 'install', '-e', '.'])
-    os.chdir('..')
+    os.system('git clone https://github.com/ZhengPeng7/BiRefNet.git')
+    sys.path.insert(0, 'BiRefNet')
+
+# Download model weights with progress bar
+def download_file(url, filename):
+    response = requests.get(url, stream=True)
+    total_size = int(response.headers.get('content-length', 0))
+    
+    with open(filename, 'wb') as f, tqdm(
+        desc=filename,
+        total=total_size,
+        unit='iB',
+        unit_scale=True,
+        unit_divisor=1024,
+    ) as bar:
+        for data in response.iter_content(chunk_size=1024):
+            size = f.write(data)
+            bar.update(size)
 
-# Download model weights
 if not os.path.exists('BiRefNet.pth'):
-    hf_hub_download(repo_id="ZhengPeng7/BiRefNet", filename="BiRefNet.pth", local_dir=".")
+    print("Downloading model weights...")
+    download_file(
+        "https://github.com/ZhengPeng7/BiRefNet/releases/download/v1.0/BiRefNet.pth",
+        "BiRefNet.pth"
+    )
 
-# Import after installation
-from BiRefNet.models.BiRefNet import BiRefNet
-from BiRefNet.utils.dataloader import test_dataset
+# Import model after setting up
+from models.BiRefNet import BiRefNet
+from utils.dataloader import test_dataset
 
-# Rest of your code remains the same...
-model = BiRefNet()
-model.load_state_dict(torch.load('BiRefNet.pth', map_location=torch.device('cpu')))
+# Initialize model
+device = torch.device('cpu')
+model = BiRefNet().to(device)
+model.load_state_dict(torch.load('BiRefNet.pth', map_location=device))
 model.eval()
 
-def remove_background(input_image):
-    # Your existing implementation
-    ...
+def process_image(input_image):
+    # Convert to numpy array
+    image = np.array(input_image)
+    original_size = image.shape[:2]
+    
+    # Resize for CPU processing
+    processed_size = (320, 320)  # Reduced size for CPU
+    image = cv2.resize(image, processed_size)
+    
+    # Preprocess
+    image = test_dataset.preprocess(image)
+    image = torch.from_numpy(image).unsqueeze(0).to(device)
+    
+    # Predict
+    with torch.no_grad():
+        pred = model(image)
+    
+    # Post-process
+    mask = (pred.squeeze().cpu().numpy() > 0.5).astype(np.uint8) * 255
+    mask = cv2.resize(mask, original_size[::-1])  # Resize back to original
+    
+    # Apply mask
+    result = cv2.bitwise_and(
+        np.array(input_image),
+        np.array(input_image),
+        mask=mask
+    )
+    
+    return Image.fromarray(result), Image.fromarray(mask)
+
+# Gradio interface
+with gr.Blocks() as demo:
+    gr.Markdown("# BiRefNet Background Remover (CPU)")
+    gr.Markdown("Works on CPU but may be slow (10-30 seconds per image)")
+    
+    with gr.Row():
+        with gr.Column():
+            input_image = gr.Image(type="pil", label="Input Image")
+            submit = gr.Button("Remove Background")
+        
+        with gr.Column():
+            output_image = gr.Image(type="pil", label="Result")
+            output_mask = gr.Image(type="pil", label="Mask")
+    
+    submit.click(
+        fn=process_image,
+        inputs=input_image,
+        outputs=[output_image, output_mask]
+    )
 
-interface = gr.Interface(...)
-interface.launch()
+demo.launch()