from ultralytics import YOLO

def train_model_and_store(model_name, epochs, img_size, batch_size, device, optimizer, learning_rate):
    # Load the pretrained model
    model = YOLO(f'./base_models/{model_name}')
    
    model.train(
        data='./BracU_Ds/data.yaml',
        epochs=epochs,
        imgsz=img_size,
        batch=batch_size,
        device=device,
        project='Activity_Detection_New_BracU',
        name=model_name,
        optimizer=optimizer,
        lr0=learning_rate,
        patience=10,
        plots=True,
        seed=42,
        # pretrained=True,
        # degrees=15,
        # translate=0.2,
        # scale=0.8,
        # shear=10.0,
        # perspective=0.001,
        # fliplr=0.5,
        # mosaic=1.0,
        # mixup=0.2,
        # hsv_h=0.015,
        # hsv_s=0.7,
        # hsv_v=0.4
    )
    
    # Load the best model after training
    # best_model_path = f'./Activity_Detection_New_BracU/{model_name}/weights/best.pt'
    # model = YOLO(best_model_path)
    
    # Evaluate the model
    # model.val(project='Activity_Detection_New_Eval', name=model_name)

if __name__ == "__main__":
    train_model_and_store(
        model_name="yolov8n.pt",
        epochs=50,
        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_BracU_Ds.txt 2>&1 &
# 2980863