File size: 3,026 Bytes
ac4f80f | 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 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 | import argparse
import subprocess
import sys
from pathlib import Path
def find_latest_best_weight(root: Path) -> Path:
train_dir = root
candidates = sorted(
train_dir.glob("run/train/weights/best.pt"),
key=lambda p: p.stat().st_mtime,
reverse=True,
)
if not candidates:
raise FileNotFoundError("No best.pt found under run/train/weights/")
return candidates[0]
def main() -> int:
parser = argparse.ArgumentParser(
description="Detect graph bounding boxes in testing folder using trained YOLOv7 model."
)
parser.add_argument(
"--weights",
type=str,
default="",
help="Optional path to model weights. If empty, latest run/train/exp5_full_finetune3/weights/best.pt is used.",
)
parser.add_argument(
"--source",
type=str,
default="testing",
help="Input image folder or file path.",
)
parser.add_argument("--img-size", type=int, default=640, help="Inference image size")
parser.add_argument("--conf-thres", type=float, default=0.45, help="Confidence threshold")
parser.add_argument("--iou-thres", type=float, default=0.45, help="NMS IoU threshold")
parser.add_argument("--device", type=str, default="0", help="Device, e.g. 0 or cpu")
parser.add_argument("--name", type=str, default="testing_graph_bbox", help="Output run name")
parser.add_argument("--view-img", action="store_true", help="Show images while running")
parser.add_argument("--no-save-txt", action="store_true", help="Do not save YOLO txt labels")
args = parser.parse_args()
root = Path(__file__).resolve().parent
source = Path(args.source)
if not source.is_absolute():
source = (root / source).resolve()
if args.weights:
weights = Path(args.weights)
if not weights.is_absolute():
weights = (root / weights).resolve()
else:
weights = find_latest_best_weight(root)
if not weights.exists():
print(f"[ERROR] Weights not found: {weights}")
return 1
if not source.exists():
print(f"[ERROR] Source not found: {source}")
return 1
cmd = [
sys.executable,
str(root / "detect.py"),
"--weights",
str(weights),
"--source",
str(source),
"--img-size",
str(args.img_size),
"--conf-thres",
str(args.conf_thres),
"--iou-thres",
str(args.iou_thres),
"--device",
args.device,
"--classes",
"0",
"--project",
"run/detect",
"--name",
args.name,
"--exist-ok",
]
if not args.no_save_txt:
cmd.extend(["--save-txt", "--save-conf"])
if args.view_img:
cmd.append("--view-img")
print("Using weights:", weights)
print("Source:", source)
print("Running command:")
print(" ".join(cmd))
return subprocess.call(cmd, cwd=str(root))
if __name__ == "__main__":
raise SystemExit(main())
|