from ultralytics import YOLO
def train_model():
    model = YOLO("yolov8-BiFPN.yaml").load('yolov8.pt')
    model.train(data="rail_defects.yaml", epochs=300)

if __name__ == '__main__':
    train_model()