File size: 1,350 Bytes
966d9af | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 | from __future__ import annotations
import argparse
from pathlib import Path
from ultralytics import YOLO
from models import register_ultralytics_modules
ROOT = Path(__file__).resolve().parent
DEFAULT_WEIGHTS = ROOT / "weights" / "symbolic_capsule_network_segmentation.pt"
def build_parser() -> argparse.ArgumentParser:
parser = argparse.ArgumentParser(description="Run Symbolic Capsule Network segmentation inference.")
parser.add_argument("source", help="Image, directory, video, or glob pattern.")
parser.add_argument("--weights", default=str(DEFAULT_WEIGHTS), help="Checkpoint path.")
parser.add_argument("--imgsz", type=int, default=640)
parser.add_argument("--conf", type=float, default=0.25)
parser.add_argument("--device", default="")
parser.add_argument("--save", action="store_true", default=True)
parser.add_argument("--show", action="store_true")
return parser
def main() -> None:
args = build_parser().parse_args()
weights = Path(args.weights).expanduser().resolve()
if not weights.exists():
raise FileNotFoundError(f"Checkpoint not found: {weights}")
register_ultralytics_modules()
model = YOLO(str(weights))
predict_kwargs = {k: v for k, v in vars(args).items() if k != "weights"}
model.predict(**predict_kwargs)
if __name__ == "__main__":
main()
|