from ultralytics import YOLO

def train_model_and_store(model_name, epochs, img_size, batch_size, device, optimizer, learning_rate):
    model = YOLO(f'./base_models/{model_name}')
    model.train(data='data/data.yaml', epochs=epochs, imgsz=img_size, batch=batch_size, device=device, project='Person_Detection', name=model_name, optimizer=optimizer, lr0=learning_rate, patience=5, plots=True, seed =42, pretrained=True)
    
    model = YOLO(f'./Activity_Detection/{model_name}/weights/best.pt')
    model.val(project='Person_Detection_Eval', name=model_name)
    

if __name__ == "__main__":
    train_model_and_store(model_name="yolov8n.pt", epochs=30, img_size=640, batch_size=64, device=0, optimizer='Adam', learning_rate=0.005)
    
# We should fine tune the following model
# yolo11n.pt
# yolov10n.pt
# yolov8n.pt
# yolov8m.pt
# yolo11n-cls.pt
# yolov8n-cls.pt

# nohup python train.py > logs/training_log_yolov8n.txt 2>&1 &
# 2083704