import yaml
from ultralytics import YOLO


yaml_path = '/content/dataset/data.yaml' 


try:
    with open(yaml_path, 'r') as f:
        data = yaml.safe_load(f)

    data['train'] = '/content/dataset/train/images'
    data['val'] = '/content/dataset/valid/images'
    
    if 'test' in data:
        data['test'] = '/content/dataset/test/images'

    with open(yaml_path, 'w') as f:
        yaml.dump(data, f)
    print("data.yaml paths successfully updated for Colab!")
    
except FileNotFoundError:
    print(f" ERROR: Could not find {yaml_path}. Please double check the unzipped folder!")

# --- 3. WAKE UP THE MODEL AND TRAIN ---
print(" Launching the Large YOLO model...")
model = YOLO("yolo26l.pt") 

results = model.train(
    data=yaml_path,      
    epochs=100,
    imgsz=640,
    batch=8,           
    project="/content/drive/MyDrive/Goyam_Colab_Training", 
    name="yolo26_large_run"             
)