File size: 1,612 Bytes
73c5503 | 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 | import argparse
import shutil
from pathlib import Path
def main():
parser = argparse.ArgumentParser(description="Treino do detector de componentes de arquitetura")
parser.add_argument("--data", type=str, default="../data/data.yaml",
help="Caminho para o data.yaml do dataset")
parser.add_argument("--epochs", type=int, default=100)
parser.add_argument("--imgsz", type=int, default=960,
help="Resolução de treino. 960 (acima do padrão 640) para preservar "
"ícones pequenos em diagramas densos.")
parser.add_argument("--model", type=str, default="yolov8n.pt",
help="Checkpoint base do YOLOv8 (transfer learning)")
parser.add_argument("--batch", type=int, default=8,
help="Batch reduzido para acomodar imgsz=960 em GPUs/CPUs limitadas.")
parser.add_argument("--out", type=str, default="../models/best.pt")
args = parser.parse_args()
from ultralytics import YOLO
model = YOLO(args.model)
results = model.train(
data=args.data,
epochs=args.epochs,
imgsz=args.imgsz,
batch=args.batch,
project="runs_threat_modeling",
name="detector_componentes",
exist_ok=True,
)
run_dir = Path(results.save_dir)
best_ckpt = run_dir / "weights" / "best.pt"
out_path = Path(args.out)
out_path.parent.mkdir(parents=True, exist_ok=True)
shutil.copy(best_ckpt, out_path)
print(f"[train] Melhor checkpoint copiado para {out_path}")
if __name__ == "__main__":
main()
|