diff --git a/backend/process.py b/backend/process.py new file mode 100644 index 0000000..4d2529a --- /dev/null +++ b/backend/process.py @@ -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() diff --git a/backend/requirements.txt b/backend/requirements.txt new file mode 100644 index 0000000..922db67 --- /dev/null +++ b/backend/requirements.txt @@ -0,0 +1,8 @@ +torch +torchvision +transformers<5 +timm +einops +kornia +pillow +accelerate diff --git a/backend/setup.py b/backend/setup.py new file mode 100644 index 0000000..be33059 --- /dev/null +++ b/backend/setup.py @@ -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()