83 lines
2.7 KiB
Python
83 lines
2.7 KiB
Python
from __future__ import annotations
|
|
|
|
import argparse
|
|
import json
|
|
import os
|
|
from pathlib import Path
|
|
|
|
from ultralytics import YOLO
|
|
|
|
DEFAULT_PROJECT_ROOT = Path(__file__).resolve().parents[4]
|
|
|
|
|
|
def project_root() -> Path:
|
|
raw = os.getenv("SEG_DATA_SERVER_ROOT")
|
|
if not raw:
|
|
return DEFAULT_PROJECT_ROOT
|
|
path = Path(raw).expanduser()
|
|
if path.is_absolute():
|
|
return path.resolve()
|
|
return (DEFAULT_PROJECT_ROOT / path).resolve()
|
|
|
|
|
|
PROJECT_ROOT = project_root()
|
|
|
|
|
|
def resolve_project_path(value: str | Path) -> Path:
|
|
path = Path(value).expanduser()
|
|
if path.is_absolute():
|
|
return path.resolve()
|
|
return (PROJECT_ROOT / path).resolve()
|
|
|
|
|
|
def main() -> None:
|
|
parser = argparse.ArgumentParser(description="Predict segmentation masks with a supplied YOLO checkpoint.")
|
|
parser.add_argument("--weights", required=True, help="Path to best.pt/last.pt or another YOLO segmentation checkpoint.")
|
|
parser.add_argument("--source", required=True, help="Image file or image directory to predict.")
|
|
parser.add_argument("--project", default=str(PROJECT_ROOT / "var" / "custom_yolo_runs"))
|
|
parser.add_argument("--name", default="custom_predict")
|
|
parser.add_argument("--imgsz", type=int, default=640)
|
|
parser.add_argument("--conf", type=float, default=0.25)
|
|
parser.add_argument("--device", default="cpu")
|
|
parser.add_argument("--save-txt", action="store_true")
|
|
parser.add_argument("--save-conf", action="store_true")
|
|
parser.add_argument("--exist-ok", action="store_true")
|
|
args = parser.parse_args()
|
|
|
|
weights = resolve_project_path(args.weights)
|
|
source = resolve_project_path(args.source)
|
|
project = resolve_project_path(args.project)
|
|
if not weights.exists():
|
|
raise FileNotFoundError(f"weights not found: {weights}")
|
|
if not source.exists():
|
|
raise FileNotFoundError(f"source not found: {source}")
|
|
|
|
model = YOLO(str(weights))
|
|
results = model.predict(
|
|
source=str(source),
|
|
imgsz=args.imgsz,
|
|
conf=args.conf,
|
|
device=args.device,
|
|
project=str(project),
|
|
name=args.name,
|
|
save=True,
|
|
save_txt=args.save_txt,
|
|
save_conf=args.save_conf,
|
|
exist_ok=args.exist_ok,
|
|
verbose=True,
|
|
)
|
|
save_dir = Path(getattr(model.predictor, "save_dir", project / args.name))
|
|
metadata = {
|
|
"weights": str(weights),
|
|
"source": str(source),
|
|
"save_dir": str(save_dir),
|
|
"count": len(results),
|
|
}
|
|
save_dir.mkdir(parents=True, exist_ok=True)
|
|
(save_dir / "predict_manifest.json").write_text(json.dumps(metadata, ensure_ascii=False, indent=2), encoding="utf-8")
|
|
print(json.dumps(metadata, ensure_ascii=False))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|