from ultralytics import YOLO

if __name__ == "__main__":
    # Load YOLO model
    model = YOLO("yolov8s.yaml")

    # Train the model
    results = model.train(
        data="datasets/data.yaml",
        epochs=100,
        imgsz=640,
        workers=0,
        batch=16,
        device="cuda",
        plots = False
    )
    # Save model
    model_path = "runs/detect/train/weights/best.pt"
    model.export(format="torchscript")
    print(f" Model saved at: {model_path}")