import os
import yaml
from ultralytics import YOLO
import streamlit as st

def create_dataset_yaml():
    """Create dataset.yaml configuration."""
    data_yaml = {
        'path': 'dataset',
        'train': 'images/train',
        'val': 'images/val',
        'names': {
            0: 'Commercial Airliner',
            1: 'Military Fighter',
            2: 'Military Transport',
            3: 'Private Jet',
            4: 'Helicopter',
            5: 'Cargo Aircraft',
            6: 'Unknown'
        }
    }
    os.makedirs('dataset/images/train', exist_ok=True)
    os.makedirs('dataset/images/val', exist_ok=True)
    os.makedirs('dataset/labels/train', exist_ok=True)
    os.makedirs('dataset/labels/val', exist_ok=True)

    with open('dataset.yaml', 'w') as f:
        yaml.dump(data_yaml, f)

def train_model():
    """Train the custom model."""
    try:
        create_dataset_yaml()
        model = YOLO('yolov8n.pt')
        model.train(
            data='dataset.yaml',
            epochs=100,
            imgsz=640,
            batch=16,
            name='aircraft_detection'
        )
        model.export(format='onnx')
        return True
    except Exception as e:
        st.error(f"Error during training: {str(e)}")
        return False