Ümit Tunç
107 lines
3.7 KiB
Python
107 lines
3.7 KiB
Python
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()
|