boc/iom/ai_pipeline/train_models.py

"""
Train AI Models
Complete training pipeline with synthetic data
"""

import sys
sys.path.insert(0, '/home/bernt/.openclaw/workspace/iom')

from ai_pipeline.data_collection import DataCollector, SyntheticDataGenerator
from ai_pipeline.training_pipeline import TrainingPipeline, TrainingConfig
import os


def create_training_dataset():
    """Create training dataset with synthetic data"""
    print("=== Creating Training Dataset ===\n")
    
    # Create collector
    collector = DataCollector("/tmp/iom_training_data")
    generator = SyntheticDataGenerator()
    
    # Generate training data
    scenes = ["street_view", "building_facade", "bridge", "road", "sidewalk", "park"]
    
    print("Generating training images...")
    for i in range(50):
        scene = scenes[i % len(scenes)]
        synthetic = generator.generate_synthetic_image(scene, num_defects=3)
        
        collector.add_annotation(
            image_id=f"train_{i:04d}",
            filename=f"train_{i:04d}.jpg",
            width=synthetic["width"],
            height=synthetic["height"],
            objects=synthetic["objects"],
            scene_type=scene,
            split="train"
        )
    
    print("Generating validation images...")
    for i in range(10):
        scene = scenes[i % len(scenes)]
        synthetic = generator.generate_synthetic_image(scene, num_defects=2)
        
        collector.add_annotation(
            image_id=f"val_{i:04d}",
            filename=f"val_{i:04d}.jpg",
            width=synthetic["width"],
            height=synthetic["height"],
            objects=synthetic["objects"],
            scene_type=scene,
            split="val"
        )
    
    print("Generating test images...")
    for i in range(10):
        scene = scenes[i % len(scenes)]
        synthetic = generator.generate_synthetic_image(scene, num_defects=2)
        
        collector.add_annotation(
            image_id=f"test_{i:04d}",
            filename=f"test_{i:04d}.jpg",
            width=synthetic["width"],
            height=synthetic["height"],
            objects=synthetic["objects"],
            scene_type=scene,
            split="test"
        )
    
    # Create data.yaml
    collector.create_data_yaml()
    
    # Stats
    stats = collector.get_stats()
    print(f"\nDataset created:")
    print(f"  Train: {stats['splits']['train']} images")
    print(f"  Val: {stats['splits']['val']} images")
    print(f"  Test: {stats['splits']['test']} images")
    print(f"  Total objects: {stats['total_objects']}")
    
    return collector


def train_yolo_model():
    """Train YOLO model"""
    print("\n=== Training YOLO Model ===\n")
    
    from ultralytics import YOLO
    
    # Load pretrained model
    model = YOLO("yolov8n.pt")
    
    # Train on synthetic data
    print("Training YOLOv8n on synthetic data...")
    results = model.train(
        data="/tmp/iom_training_data/data.yaml",
        epochs=5,  # Reduced for demo
        batch=8,
        imgsz=640,
        device="cpu",
        project="/tmp/iom_models",
        name="yolo_infrastructure",
        exist_ok=True
    )
    
    print(f"Training complete!")
    print(f"Model saved: /tmp/iom_models/yolo_infrastructure/weights/best.pt")
    
    return model


def evaluate_model(model):
    """Evaluate trained model"""
    print("\n=== Evaluating Model ===\n")
    
    # Validate on test set
    metrics = model.val()
    
    print("Evaluation results:")
    print(f"  mAP50: {metrics.box.map50:.3f}")
    print(f"  mAP50-95: {metrics.box.map:.3f}")
    print(f"  Precision: {metrics.box.mp:.3f}")
    print(f"  Recall: {metrics.box.mr:.3f}")
    
    return metrics


def export_model(model):
    """Export model to production format"""
    print("\n=== Exporting Model ===\n")
    
    # Export to ONNX
    print("Exporting to ONNX...")
    model.export(format="onnx", dynamic=True)
    
    # Export to TorchScript
    print("Exporting to TorchScript...")
    model.export(format="torchscript")
    
    print("Export complete!")
    print("Formats: PyTorch, ONNX, TorchScript")


def main():
    """Main training pipeline"""
    print("=" * 60)
    print("IOM AI MODEL TRAINING")
    print("=" * 60)
    
    # 1. Create dataset
    dataset = create_training_dataset()
    
    # 2. Train model
    model = train_yolo_model()
    
    # 3. Evaluate
    metrics = evaluate_model(model)
    
    # 4. Export
    export_model(model)
    
    print("\n" + "=" * 60)
    print("TRAINING COMPLETE")
    print("=" * 60)
    print("\nModels saved to: /tmp/iom_models/")
    print("Dataset saved to: /tmp/iom_training_data/")
    print("\nNext steps:")
    print("1. Collect real infrastructure images")
    print("2. Annotate with Label Studio")
    print("3. Retrain with real data")
    print("4. Deploy to production")


if __name__ == '__main__':
    main()