from ultralytics import YOLO

def train_model():
    # Load YOLOv8n (nano) for faster training with decent accuracy
    model = YOLO('yolov8n.pt')
    
    # Train with optimized parameters for speed and quality
    results = model.train(
        data='dataset.yaml',
        epochs=50,           # Reduced number of epochs
        imgsz=640,           # Standard image size for faster processing
        batch=8,             # Smaller batch size for less memory usage
        name='memory_detector_fast',
        save=True,
        device='cpu',
        patience=15,         # Shorter patience for earlier stopping
        save_period=5,       # Save every 5 epochs
        verbose=True,
        
        # Effective but lightweight augmentation
        degrees=5.0,         # Less rotation for speed
        scale=0.5,
        translate=0.1,
        fliplr=0.5,
        mosaic=1.0,          # Keep mosaic as it's very effective
        
        # Speed-optimized optimization parameters
        lr0=0.01,
        lrf=0.01,
        momentum=0.937,
        weight_decay=0.0005,
        warmup_epochs=1.0,   # Shorter warmup
        
        # Performance parameters
        workers=0,           # Fewer workers for CPU training
        cache='disk',        # Changed to disk caching for deterministic results
    )
    
     # Create model directory if it doesn't exist
    import os
    os.makedirs('model/weights', exist_ok=True)

    # Save the trained model
    model.save('model/weights/best.pt')


if __name__ == '__main__':
    train_model()