91 lines
3.5 KiB
Python
91 lines
3.5 KiB
Python
#!/usr/bin/env python3
|
|
"""FOBG — AI Background Remover — Entry Point.
|
|
|
|
Usage:
|
|
python main.py # Launch Gradio UI (default)
|
|
python main.py --host 0.0.0.0 --port 7860
|
|
python main.py --share # Create public Gradio share link
|
|
python main.py --cli --help # CLI mode (scripted use)
|
|
"""
|
|
|
|
import argparse
|
|
import sys
|
|
import os
|
|
|
|
# Ensure the project root is on sys.path
|
|
_PROJECT_ROOT = os.path.dirname(os.path.abspath(__file__))
|
|
sys.path.insert(0, _PROJECT_ROOT)
|
|
|
|
from src.ui.app import launch_app
|
|
|
|
|
|
def cli_demo():
|
|
"""Simple CLI demo for batch image processing."""
|
|
import argparse as ap2
|
|
parser = ap2.ArgumentParser(description="FOBG CLI — process images/videos")
|
|
parser.add_argument("input", help="Input image or video file path")
|
|
parser.add_argument("-o", "--output", help="Output file path (auto-generated if omitted)")
|
|
parser.add_argument("--model", default="u2net", choices=["u2net", "u2netp", "u2net_human_seg", "isnet-general-use", "isnet-anime"], help="rembg model")
|
|
parser.add_argument("--alpha-matting", action="store_true", default=True, help="Use alpha matting")
|
|
parser.add_argument("--no-alpha-matting", action="store_true", help="Disable alpha matting")
|
|
parser.add_argument("--video", action="store_true", help="Treat input as video (frame-by-frame processing)")
|
|
args = parser.parse_args()
|
|
|
|
input_path = args.input
|
|
if not os.path.exists(input_path):
|
|
print(f"Error: File not found: {input_path}")
|
|
sys.exit(1)
|
|
|
|
is_video = args.video or input_path.lower().endswith((".mp4", ".avi", ".webm", ".mov", ".mkv"))
|
|
output = args.output or f"output_{os.path.splitext(os.path.basename(input_path))[0]}_{os.urandom(4).hex()}"
|
|
if is_video:
|
|
output += ".mp4"
|
|
else:
|
|
output += ".png"
|
|
|
|
from src.core.bg_remove import remove_background_image, remove_background_bytes
|
|
from src.core.video_bg_remove import process_video_frame_by_frame
|
|
|
|
if is_video:
|
|
print(f"Processing video: {input_path}")
|
|
result = process_video_frame_by_frame(
|
|
input_path, output,
|
|
model=args.model,
|
|
alpha_matting=not args.no_alpha_matting,
|
|
)
|
|
print(f"Video saved to: {result}")
|
|
else:
|
|
print(f"Processing image: {input_path}")
|
|
from PIL import Image
|
|
input_img = Image.open(input_path).convert("RGB")
|
|
result = remove_background_image(input_img, model=args.model, force=True, alpha_matting=not args.no_alpha_matting)
|
|
result.save(output, format="PNG")
|
|
print(f"Image saved to: {output}")
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(description="FOBG — AI Background Remover")
|
|
parser.add_argument("--cli", action="store_true", help="Run in CLI mode for batch processing")
|
|
parser.add_argument("--host", default="0.0.0.0", help="Gradio server host")
|
|
parser.add_argument("--port", type=int, default=7860, help="Gradio server port")
|
|
parser.add_argument("--share", action="store_true", help="Create a public share link")
|
|
parser.add_argument("--no-browser", action="store_true", help="Do not open browser")
|
|
|
|
args = parser.parse_args()
|
|
|
|
if args.cli:
|
|
# Parse remaining args for CLI mode
|
|
sys.argv = ["fobg"] + [arg for arg in sys.argv[1:] if arg not in ("--cli",)]
|
|
cli_demo()
|
|
else:
|
|
launch_app(
|
|
host=args.host,
|
|
port=args.port,
|
|
share=args.share,
|
|
inbrowser=not args.no_browser,
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|