advanced_train.py

#!/usr/bin/env python3
"""
Advanced YOLO11 training cho Night Market People Detection
với custom augmentations tối ưu cho điều kiện ánh sáng yếu
"""

import torch
import torch.nn as nn
import torchvision.transforms as transforms
from ultralytics import YOLO
from ultralytics.data.augment import Compose, Albumentations
import albumentations as A
import cv2
import numpy as np
import yaml
import os
from pathlib import Path

class NightSceneAugmentation:
    """Custom augmentation pipeline cho phát hiện người trong cảnh đêm"""
    
    def __init__(self, p=0.5):
        self.augment = A.Compose([
            # Augmentations đặc biệt cho ánh sáng yếu
            A.RandomBrightnessContrast(
                brightness_limit=(-0.3, 0.3),
                contrast_limit=(-0.3, 0.3),
                p=0.8
            ),
            A.CLAHE(clip_limit=4.0, tile_grid_size=(8, 8), p=0.5),
            A.RandomGamma(gamma_limit=(70, 130), p=0.5),
            A.HueSaturationValue(
                hue_shift_limit=20,
                sat_shift_limit=30,
                val_shift_limit=30,
                p=0.7
            ),
            
            # Noise augmentations cho cảnh đêm
            A.OneOf([
                A.GaussNoise(var_limit=(10.0, 50.0), p=0.3),
                A.ISONoise(color_shift=(0.01, 0.05), intensity=(0.1, 0.5), p=0.3),
                A.MultiplicativeNoise(multiplier=(0.9, 1.1), p=0.3),
            ], p=0.4),
            
            # Hiệu ứng thời tiết và khí quyển
            A.OneOf([
                A.RandomFog(fog_coef_lower=0.1, fog_coef_upper=0.3, p=0.2),
                A.RandomRain(
                    slant_lower=-10, slant_upper=10,
                    drop_length=20, drop_width=1,
                    drop_color=(200, 200, 200),
                    blur_value=1, p=0.1
                ),
            ], p=0.2),
            
            # Biến đổi hình học
            A.ShiftScaleRotate(
                shift_limit=0.1,
                scale_limit=0.2,
                rotate_limit=10,
                border_mode=cv2.BORDER_REFLECT_101,
                p=0.6
            ),
            A.HorizontalFlip(p=0.5),
            
            # Hiệu ứng ống kính phổ biến trong camera giám sát
            A.OneOf([
                A.MotionBlur(blur_limit=3, p=0.3),
                A.Defocus(radius=(1, 3), alias_blur=(0.1, 0.5), p=0.2),
                A.ZoomBlur(max_factor=1.05, p=0.1),
            ], p=0.3),
            
        ], bbox_params=A.BboxParams(
            format='yolo',
            label_fields=['class_labels'],
            min_visibility=0.3
        ))
    
    def __call__(self, image, bboxes, class_labels):
        """Áp dụng augmentations cho image và bounding boxes"""
        try:
            augmented = self.augment(
                image=image,
                bboxes=bboxes,
                class_labels=class_labels
            )
            return augmented['image'], augmented['bboxes'], augmented['class_labels']
        except:
            # Trả về original nếu augmentation thất bại
            return image, bboxes, class_labels

def create_enhanced_config():
    """Tạo enhanced training configuration"""
    config = {
        # Dataset
        'data': 'Night-Market-People-Detection-14/data.yaml',
        
        # Model
        'model': 'yolo11m.pt',
        
        # Tham số training
        'epochs': 150,  # Tăng để hội tụ tốt hơn
        'batch': 16,
        'imgsz': 800,
        'device': 'cuda' if torch.cuda.is_available() else 'cpu',
        'workers': 8,
        'patience': 20,
        'save_period': 10,
        
        # Output
        'project': 'night_market_detection',
        'name': 'enhanced_yolo11m_night',
        
        # Tối ưu hóa
        'optimizer': 'AdamW',
        'lr0': 0.0008,  # Learning rate ban đầu thấp hơn một chút
        'lrf': 0.01,
        'momentum': 0.937,
        'weight_decay': 0.0005,
        'warmup_epochs': 5.0,
        'warmup_momentum': 0.8,
        'warmup_bias_lr': 0.1,
        
        # Trọng số loss (điều chỉnh cho phát hiện người)
        'box': 7.5,
        'cls': 0.8,    # Trọng số classification cao hơn
        'dfl': 1.5,
        
        # Cài đặt training nâng cao
        'amp': True,
        'deterministic': True,
        'seed': 42,
        'resume': False,
        'pretrained': True,
        
        # Data augmentation (tối ưu cho đêm)
        'augment': True,
        'hsv_h': 0.02,   # Biến thiên hue cao hơn một chút
        'hsv_s': 0.8,    # Biến thiên saturation cao hơn cho cảnh đêm
        'hsv_v': 0.5,    # Biến thiên value cao hơn cho thay đổi ánh sáng
        'degrees': 8.0,  # Xoay vừa phải cho góc giám sát
        'translate': 0.1,
        'scale': 0.6,    # Biến thiên scale cao hơn
        'shear': 3.0,
        'perspective': 0.0001,
        'flipud': 0.0,
        'fliplr': 0.5,
        'mosaic': 1.0,
        'mixup': 0.15,   # Tăng mixup để tổng quát hóa tốt hơn
        'copy_paste': 0.1,
        'auto_augment': 'randaugment',
        'erasing': 0.4,
        'crop_fraction': 1.0,
        'close_mosaic': 15,
        
        # Validation
        'val': True,
        'split': 'val',
        'save_json': True,
        'save_hybrid': False,
        'conf': 0.001,   # Conf thấp hơn cho validation
        'iou': 0.7,
        'max_det': 300,
        'half': False,
        'dnn': False,
        'plots': True,
        'source': None,
        'show': False,
        'save_txt': False,
        'save_conf': False,
        'save_crop': False,
        'show_labels': True,
        'show_conf': True,
        'vid_stride': 1,
        'line_width': None,
        'visualize': False,
        # 'augment_val': False,  # Not a valid YOLO argument
        'agnostic_nms': False,
        'retina_masks': False,
        'classes': None,
        'verbose': True,
    }
    
    return config

def train_with_callbacks():
    """Train YOLO11 với custom callbacks để monitoring"""
    
    def on_train_start(trainer):
        print("🚀 Bắt đầu enhanced night market people detection training")
        print(f"📊 Dataset: {trainer.args.data}")
        print(f"🏃‍♂️ Epochs: {trainer.args.epochs}")
        print(f"📦 Batch size: {trainer.args.batch}")
        print(f"🖼️  Image size: {trainer.args.imgsz}")
        print(f"💻 Device: {trainer.args.device}")
        
    def on_epoch_end(trainer):
        if trainer.epoch % 10 == 0:
            print(f"📈 Epoch {trainer.epoch}: Loss={trainer.loss:.4f}")
            
    def on_train_end(trainer):
        print("✅ Training hoàn thành thành công!")
        print(f"💾 Best model: {trainer.best}")
        print(f"📁 Results: {trainer.save_dir}")
    
    # Load model
    model = YOLO('yolo11m.pt')
    
    # Thêm callbacks
    model.add_callback('on_train_start', on_train_start)
    model.add_callback('on_epoch_end', on_epoch_end)
    model.add_callback('on_train_end', on_train_end)
    
    # Lấy enhanced configuration
    config = create_enhanced_config()
    
    print("🔧 Cấu hình Training:")
    for key, value in config.items():
        if key not in ['data']:
            print(f"   {key}: {value}")
    
    # Bắt đầu training
    results = model.train(**config)
    
    return model, results

def main():
    """Hàm training chính"""
    print("🌙 Enhanced YOLO11 Training cho Night Market People Detection")
    print("=" * 60)
    
    # Kiểm tra CUDA có sẵn
    if torch.cuda.is_available():
        print(f"🎮 CUDA có sẵn: {torch.cuda.get_device_name()}")
        print(f"💾 GPU Memory: {torch.cuda.get_device_properties(0).total_memory / 1e9:.1f} GB")
    else:
        print("⚠️  CUDA không có sẵn, training trên CPU")
    
    # Train model
    model, results = train_with_callbacks()
    
    # Validation
    print("\n🔍 Đang chạy validation...")
    val_results = model.val(
        data='Night-Market-People-Detection-14/data.yaml',
        imgsz=800,
        conf=0.25,
        iou=0.7,
        split='val'
    )
    
    print(f"\n📊 Kết quả Validation:")
    print(f"   📈 mAP50: {val_results.box.map50:.4f}")
    print(f"   📈 mAP50-95: {val_results.box.map:.4f}")
    print(f"   📈 Precision: {val_results.box.mp:.4f}")
    print(f"   📈 Recall: {val_results.box.mr:.4f}")
    
    # Test evaluation
    print("\n🧪 Đang chạy test evaluation...")
    test_results = model.val(
        data='Night-Market-People-Detection-14/data.yaml',
        imgsz=800,
        conf=0.25,
        iou=0.7,
        split='test'
    )
    
    print(f"\n🎯 Kết quả Test:")
    print(f"   📈 mAP50: {test_results.box.map50:.4f}")
    print(f"   📈 mAP50-95: {test_results.box.map:.4f}")
    print(f"   📈 Precision: {test_results.box.mp:.4f}")
    print(f"   📈 Recall: {test_results.box.mr:.4f}")
    
    # Export optimized model
    print("\n📦 Đang export optimized model...")
    model.export(
        format='torchscript',
        optimize=True,
        half=False,
        int8=False,
        dynamic=False,
        simplify=True
    )
    
    print("\n✨ Training pipeline hoàn thành thành công!")
    return model, results

if __name__ == "__main__":
    main()