thawro/yolov8-digits-detection

1from ultralytics import YOLO2 3 4def train(5    ckpt_path: str = "yolov8n.pt",6    epochs: int = 50,7    imgsz: tuple[int, int] = (256, 256),8    batch: int = 128,9    export_format: str | None = "onnx",10):11    model = YOLO(ckpt_path)12    model.train(data="yolo_HWD+.yaml", epochs=epochs, imgsz=imgsz, device=0, batch=batch)13    if export_format is not None:14        model.export(format=export_format, imgsz=imgsz)15 16 17if __name__ == "__main__":18    IMGSZ = 25619    CKPT_PATH = "yolov8n.pt"20    EPOCHS = 5021    IMGSZ = (256, 256)22    BATCH = 12823    EXPORT_FORMAT = "onnx"24    train(ckpt_path=CKPT_PATH, epochs=EPOCHS, imgsz=IMGSZ, batch=BATCH, export_format=EXPORT_FORMAT)25