feat: implement Python BiRefNet background segmentation processor and installer

2026-05-23 10:48:41 +03:00
parent 7136b5bbd9
commit 331487549e
3 changed files with 200 additions and 0 deletions
@@ -0,0 +1,106 @@
+import sys
+import os
+import argparse
+
+# Force output buffering off so Electron gets real-time status prints
+sys.stdout.reconfigure(line_buffering=True)
+
+print("[STATUS] Importing libraries...")
+try:
+    import torch
+    from torchvision import transforms
+    from PIL import Image
+    from transformers import AutoModelForImageSegmentation
+except Exception as e:
+    print(f"[ERROR] Dependency import failed: {str(e)}", file=sys.stderr)
+    sys.exit(1)
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--input', required=True, help='Path to input image')
+    parser.add_argument('--output', required=True, help='Path to save output image')
+    parser.add_argument('--model', default='ZhengPeng7/BiRefNet', help='BiRefNet model name')
+    parser.add_argument('--device', default='auto', help='Device (cuda, cpu, auto)')
+    args = parser.parse_args()
+
+    if not os.path.exists(args.input):
+        print(f"[ERROR] Input file {args.input} does not exist", file=sys.stderr)
+        sys.exit(1)
+
+    print("[STATUS] Detecting device...")
+    if args.device == 'auto':
+        device = 'cuda' if torch.cuda.is_available() else 'cpu'
+    else:
+        device = args.device
+    print(f"[STATUS] Using device: {device.upper()}")
+
+    print("[STATUS] Loading BiRefNet model (this may take a moment on first run)...")
+    try:
+        # Load the model with trust_remote_code=True
+        model = AutoModelForImageSegmentation.from_pretrained(args.model, trust_remote_code=True)
+        model.to(device)
+        model.eval()
+    except Exception as e:
+        print(f"[ERROR] Failed to load model: {str(e)}", file=sys.stderr)
+        sys.exit(1)
+
+    print("[STATUS] Loading and preprocessing image...")
+    try:
+        orig_img = Image.open(args.input)
+        
+        # Keep original image color profile and transparency if any, but convert to RGB for model
+        if orig_img.mode != 'RGB':
+            img = orig_img.convert('RGB')
+        else:
+            img = orig_img
+
+        # BiRefNet works best with 1024x1024 input
+        transform_image = transforms.Compose([
+            transforms.Resize((1024, 1024)),
+            transforms.ToTensor(),
+            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
+        ])
+        
+        input_tensor = transform_image(img).unsqueeze(0).to(device)
+    except Exception as e:
+        print(f"[ERROR] Failed to preprocess image: {str(e)}", file=sys.stderr)
+        sys.exit(1)
+
+    print("[STATUS] Running background removal inference...")
+    try:
+        with torch.no_grad():
+            preds = model(input_tensor)[-1].sigmoid().cpu()
+        
+        pred = preds[0].squeeze()
+    except Exception as e:
+        print(f"[ERROR] Inference failed: {str(e)}", file=sys.stderr)
+        sys.exit(1)
+
+    print("[STATUS] Generating transparency mask...")
+    try:
+        # Resize mask back to original image size
+        mask = transforms.ToPILImage()(pred).resize(orig_img.size, Image.Resampling.BILINEAR)
+
+        # Convert original image to RGBA and set the transparency alpha channel
+        rgba_img = orig_img.convert("RGBA")
+        rgba_img.putalpha(mask)
+    except Exception as e:
+        print(f"[ERROR] Mask generation failed: {str(e)}", file=sys.stderr)
+        sys.exit(1)
+
+    print(f"[STATUS] Saving output transparent image...")
+    try:
+        # Create output directory if it doesn't exist
+        out_dir = os.path.dirname(args.output)
+        if out_dir and not os.path.exists(out_dir):
+            os.makedirs(out_dir)
+            
+        rgba_img.save(args.output, "PNG")
+    except Exception as e:
+        print(f"[ERROR] Failed to save output image: {str(e)}", file=sys.stderr)
+        sys.exit(1)
+
+    print("[STATUS] DONE")
+
+if __name__ == '__main__':
+    main()
@@ -0,0 +1,8 @@
+torch
+torchvision
+transformers<5
+timm
+einops
+kornia
+pillow
+accelerate
@@ -0,0 +1,86 @@
+import os
+import sys
+import subprocess
+import venv
+
+# Force line buffering
+sys.stdout.reconfigure(line_buffering=True)
+
+def run_command(command, description):
+    print(f"[STATUS] {description}...")
+    try:
+        process = subprocess.Popen(
+            command,
+            stdout=subprocess.PIPE,
+            stderr=subprocess.STDOUT,
+            text=True,
+            shell=True,
+            bufsize=1
+        )
+        
+        # Read output in real-time
+        for line in process.stdout:
+            line_str = line.strip()
+            # If the output has pip download status, keep it readable but don't spam too much
+            if "Downloading" in line_str or "Installing" in line_str:
+                print(f"[STATUS] {line_str}")
+            elif line_str:
+                print(f"[INFO] {line_str}")
+                
+        process.wait()
+        if process.returncode != 0:
+            print(f"[ERROR] Command failed with code {process.returncode}", file=sys.stderr)
+            return False
+        return True
+    except Exception as e:
+        print(f"[ERROR] Exception occurred: {str(e)}", file=sys.stderr)
+        return False
+
+def main():
+    base_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+    venv_dir = os.path.join(base_dir, ".venv")
+    
+    # 1. Create venv if not exists
+    if not os.path.exists(venv_dir):
+        print("[STATUS] Creating Python virtual environment...")
+        try:
+            venv.create(venv_dir, with_pip=True)
+            print("[STATUS] Virtual environment successfully created.")
+        except Exception as e:
+            print(f"[ERROR] Failed to create virtual environment: {str(e)}", file=sys.stderr)
+            sys.exit(1)
+    else:
+        print("[STATUS] Virtual environment already exists.")
+
+    # 2. Path to pip inside the venv
+    if os.name == 'nt':
+        pip_path = os.path.join(venv_dir, "Scripts", "pip.exe")
+        python_path = os.path.join(venv_dir, "Scripts", "python.exe")
+    else:
+        pip_path = os.path.join(venv_dir, "bin", "pip")
+        python_path = os.path.join(venv_dir, "bin", "python")
+
+    if not os.path.exists(pip_path):
+        print("[ERROR] pip executable not found inside virtual environment!", file=sys.stderr)
+        sys.exit(1)
+
+    # 3. Upgrade pip
+    run_command(f'"{pip_path}" install --upgrade pip', "Upgrading pip to latest version")
+
+    # 4. Install dependencies
+    requirements_path = os.path.join(base_dir, "backend", "requirements.txt")
+    
+    # Use extra index URL for CPU-only torch to download faster (150MB instead of 2.5GB+ GPU version)
+    # This is perfect for single images on desktop!
+    pip_install_cmd = f'"{pip_path}" install -r "{requirements_path}" --extra-index-url https://download.pytorch.org/whl/cpu'
+    
+    success = run_command(pip_install_cmd, "Installing PyTorch, Transformers, and PIL dependencies")
+    if not success:
+        print("[ERROR] Dependency installation failed!", file=sys.stderr)
+        sys.exit(1)
+
+    print("[STATUS] Setup completed successfully!")
+    print("[STATUS] DONE")
+
+if __name__ == '__main__':
+    main()