DefendIntelligence/vessel-detection

1from __future__ import annotations2 3import argparse4import os5import subprocess6import urllib.request7import venv8from pathlib import Path9 10 11ROOT = Path(__file__).resolve().parent12VENV_DIR = ROOT / ".venv"13MODEL_DIR = ROOT / "models"14MODEL_PATH = MODEL_DIR / "best.pt"15MODEL_URL = "https://huggingface.co/DefendIntelligence/vessel-detection/resolve/main/models/best.pt"16 17 18def _venv_python() -> Path:19    if os.name == "nt":20        return VENV_DIR / "Scripts" / "python.exe"21    return VENV_DIR / "bin" / "python"22 23 24def _run(command: list[str | os.PathLike[str]], env: dict[str, str] | None = None) -> None:25    printable = " ".join(str(part) for part in command)26    print(f"\n$ {printable}", flush=True)27    subprocess.check_call([str(part) for part in command], cwd=ROOT, env=env)28 29 30def _ensure_venv() -> Path:31    python_path = _venv_python()32    if not python_path.exists():33        print(f"Creating virtual environment: {VENV_DIR}", flush=True)34        venv.EnvBuilder(with_pip=True).create(VENV_DIR)35    return python_path36 37 38def _install_dependencies(python_path: Path) -> None:39    _run([python_path, "-m", "pip", "install", "--upgrade", "pip"])40    _run([python_path, "-m", "pip", "install", "-r", "requirements.txt"])41 42 43def _download_model() -> None:44    MODEL_DIR.mkdir(parents=True, exist_ok=True)45    if MODEL_PATH.exists() and MODEL_PATH.stat().st_size > 0:46        print(f"Model already present: {MODEL_PATH}", flush=True)47        return48 49    tmp_path = MODEL_PATH.with_suffix(".pt.tmp")50    print(f"Downloading model from Hugging Face:\n{MODEL_URL}", flush=True)51    with urllib.request.urlopen(MODEL_URL) as response, tmp_path.open("wb") as handle:52        total = int(response.headers.get("Content-Length") or 0)53        downloaded = 054        while True:55            chunk = response.read(1024 * 1024)56            if not chunk:57                break58            handle.write(chunk)59            downloaded += len(chunk)60            if total:61                percent = downloaded * 100 / total62                print(f"\r{downloaded / 1_000_000:.1f} MB / {total / 1_000_000:.1f} MB ({percent:.0f}%)", end="")63            else:64                print(f"\r{downloaded / 1_000_000:.1f} MB", end="")65        print()66    tmp_path.replace(MODEL_PATH)67    print(f"Saved model to: {MODEL_PATH}", flush=True)68 69 70def main() -> None:71    parser = argparse.ArgumentParser(description="Install and run the Vessel Detection Gradio demo locally.")72    parser.add_argument("--skip-install", action="store_true", help="Do not install Python dependencies.")73    parser.add_argument("--download-only", action="store_true", help="Download the model and exit.")74    parser.add_argument("--host", default="127.0.0.1", help="Gradio server host.")75    parser.add_argument("--port", default="7860", help="Gradio server port.")76    args = parser.parse_args()77 78    python_path = None79    if not (args.download_only and args.skip_install):80        python_path = _ensure_venv()81    if not args.skip_install:82        if python_path is None:83            python_path = _ensure_venv()84        _install_dependencies(python_path)85    _download_model()86 87    if args.download_only:88        return89 90    if python_path is None:91        python_path = _ensure_venv()92    env = os.environ.copy()93    env["GRADIO_SERVER_NAME"] = args.host94    env["GRADIO_SERVER_PORT"] = args.port95    print(f"\nStarting Gradio at http://{args.host}:{args.port}", flush=True)96    _run([python_path, "app.py"], env=env)97 98 99if __name__ == "__main__":100    main()101