train_model.py

"""
Training Module for Damaged Poster Detection Model
================================================

This module handles training of the deep learning model for damage classification.
Includes data augmentation, model training, validation, and model saving.
"""

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader
import torchvision.transforms as transforms
from torchvision.models import efficientnet_b0, EfficientNet_B0_Weights

import cv2
import numpy as np
from PIL import Image
import albumentations as A
from albumentations.pytorch import ToTensorV2
import pandas as pd
from sklearn.model_selection import train_test_split
from sklearn.metrics import classification_report, confusion_matrix
import matplotlib.pyplot as plt
import seaborn as sns

import os
import json
import logging
from pathlib import Path
from typing import Dict, List, Tuple, Optional
from dataclasses import dataclass
from datetime import datetime
import wandb  # For experiment tracking

from damaged_poster_detector import DamageClassifier, DamageConfig

# Setup logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


@dataclass
class TrainingConfig:
    """Configuration for model training"""
    # Data
    data_dir: str = "data"
    train_split: float = 0.8
    val_split: float = 0.15
    test_split: float = 0.05
    
    # Training
    batch_size: int = 32
    num_epochs: int = 50
    learning_rate: float = 1e-4
    weight_decay: float = 1e-5
    patience: int = 10
    
    # Model
    num_classes: int = 7
    model_name: str = "efficientnet_b0"
    pretrained: bool = True
    
    # Augmentation
    use_augmentation: bool = True
    augmentation_prob: float = 0.8
    
    # Output
    model_save_path: str = "models"
    log_dir: str = "logs"
    
    # Hardware
    device: str = "auto"
    num_workers: int = 4


class DamageDataset(Dataset):
    """Dataset class for damage detection training"""
    
    def __init__(self, 
                 image_paths: List[str], 
                 labels: List[int],
                 transform=None,
                 augment: bool = False):
        """
        Initialize dataset
        
        Args:
            image_paths: List of image file paths
            labels: List of damage labels (0=no_damage, 1=tear, 2=crease, etc.)
            transform: Image transformations
            augment: Whether to apply data augmentation
        """
        self.image_paths = image_paths
        self.labels = labels
        self.transform = transform
        self.augment = augment
        
        # Define augmentation pipeline
        if augment:
            self.augmentation = A.Compose([
                A.RandomRotate90(p=0.5),
                A.Flip(p=0.5),
                A.RandomBrightnessContrast(
                    brightness_limit=0.2, 
                    contrast_limit=0.2, 
                    p=0.5
                ),
                A.OneOf([
                    A.GaussNoise(var_limit=(10.0, 50.0), p=1.0),
                    A.GaussianBlur(blur_limit=3, p=1.0),
                ], p=0.3),
                A.OneOf([
                    A.HueSaturationValue(
                        hue_shift_limit=10,
                        sat_shift_limit=15,
                        val_shift_limit=10,
                        p=1.0
                    ),
                    A.CLAHE(clip_limit=2.0, tile_grid_size=(8, 8), p=1.0),
                ], p=0.3),
                A.CoarseDropout(
                    max_holes=8,
                    max_height=32,
                    max_width=32,
                    min_holes=1,
                    min_height=8,
                    min_width=8,
                    fill_value=0,
                    p=0.3
                ),
            ])
    
    def __len__(self):
        return len(self.image_paths)
    
    def __getitem__(self, idx):
        # Load image
        image_path = self.image_paths[idx]
        image = cv2.imread(image_path)
        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        
        # Apply augmentation if enabled
        if self.augment:
            augmented = self.augmentation(image=image)
            image = augmented['image']
        
        # Apply transforms
        if self.transform:
            image = self.transform(image)
        
        label = torch.tensor(self.labels[idx], dtype=torch.long)
        
        return image, label


class DamageTrainer:
    """Main trainer class for damage detection model"""
    
    def __init__(self, config: TrainingConfig):
        """
        Initialize trainer
        
        Args:
            config: Training configuration
        """
        self.config = config
        
        # Setup device
        if config.device == "auto":
            self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        else:
            self.device = torch.device(config.device)
        
        logger.info(f"Using device: {self.device}")
        
        # Create directories
        os.makedirs(config.model_save_path, exist_ok=True)
        os.makedirs(config.log_dir, exist_ok=True)
        
        # Initialize model
        self.model = DamageClassifier(num_classes=config.num_classes).to(self.device)
        
        # Define transforms
        self.train_transform = transforms.Compose([
            transforms.ToPILImage(),
            transforms.Resize((224, 224)),
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
        ])
        
        self.val_transform = transforms.Compose([
            transforms.ToPILImage(),
            transforms.Resize((224, 224)),
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
        ])
        
        # Training history
        self.history = {
            'train_loss': [],
            'train_acc': [],
            'val_loss': [],
            'val_acc': []
        }
    
    def prepare_data(self, data_dir: str) -> Tuple[DataLoader, DataLoader, DataLoader]:
        """
        Prepare training, validation, and test data loaders
        
        Args:
            data_dir: Directory containing training data
            
        Returns:
            Tuple of (train_loader, val_loader, test_loader)
        """
        # Expected directory structure:
        # data/
        #   ├── no_damage/
        #   ├── tear/
        #   ├── crease/
        #   ├── stain/
        #   ├── fade/
        #   ├── burn/
        #   └── water_damage/
        
        damage_types = [
            'no_damage', 'tear', 'crease', 'stain', 
            'fade', 'burn', 'water_damage'
        ]
        
        image_paths = []
        labels = []
        
        # Collect all images and labels
        for label_idx, damage_type in enumerate(damage_types):
            damage_dir = Path(data_dir) / damage_type
            if damage_dir.exists():
                for img_path in damage_dir.glob('*'):
                    if img_path.suffix.lower() in ['.jpg', '.jpeg', '.png', '.bmp']:
                        image_paths.append(str(img_path))
                        labels.append(label_idx)
        
        if not image_paths:
            raise ValueError(f"No images found in {data_dir}")
        
        logger.info(f"Found {len(image_paths)} images across {len(set(labels))} classes")
        
        # Split data
        train_paths, temp_paths, train_labels, temp_labels = train_test_split(
            image_paths, labels, 
            test_size=(self.config.val_split + self.config.test_split),
            stratify=labels, 
            random_state=42
        )
        
        val_paths, test_paths, val_labels, test_labels = train_test_split(
            temp_paths, temp_labels,
            test_size=self.config.test_split / (self.config.val_split + self.config.test_split),
            stratify=temp_labels,
            random_state=42
        )
        
        # Create datasets
        train_dataset = DamageDataset(
            train_paths, train_labels, 
            transform=self.train_transform,
            augment=self.config.use_augmentation
        )
        
        val_dataset = DamageDataset(
            val_paths, val_labels,
            transform=self.val_transform,
            augment=False
        )
        
        test_dataset = DamageDataset(
            test_paths, test_labels,
            transform=self.val_transform,
            augment=False
        )
        
        # Create data loaders
        train_loader = DataLoader(
            train_dataset,
            batch_size=self.config.batch_size,
            shuffle=True,
            num_workers=self.config.num_workers,
            pin_memory=True if self.device.type == 'cuda' else False
        )
        
        val_loader = DataLoader(
            val_dataset,
            batch_size=self.config.batch_size,
            shuffle=False,
            num_workers=self.config.num_workers,
            pin_memory=True if self.device.type == 'cuda' else False
        )
        
        test_loader = DataLoader(
            test_dataset,
            batch_size=self.config.batch_size,
            shuffle=False,
            num_workers=self.config.num_workers,
            pin_memory=True if self.device.type == 'cuda' else False
        )
        
        logger.info(f"Train: {len(train_dataset)}, Val: {len(val_dataset)}, Test: {len(test_dataset)}")
        
        return train_loader, val_loader, test_loader
    
    def train_epoch(self, train_loader: DataLoader, optimizer, criterion) -> Tuple[float, float]:
        """Train for one epoch"""
        
        self.model.train()
        running_loss = 0.0
        correct = 0
        total = 0
        
        for batch_idx, (images, labels) in enumerate(train_loader):
            images, labels = images.to(self.device), labels.to(self.device)
            
            # Zero gradients
            optimizer.zero_grad()
            
            # Forward pass
            outputs = self.model(images)
            loss = criterion(outputs, labels)
            
            # Backward pass
            loss.backward()
            optimizer.step()
            
            # Statistics
            running_loss += loss.item()
            _, predicted = torch.max(outputs.data, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()
            
            if batch_idx % 50 == 0:
                logger.info(f'Batch {batch_idx}/{len(train_loader)}, Loss: {loss.item():.4f}')
        
        epoch_loss = running_loss / len(train_loader)
        epoch_acc = correct / total
        
        return epoch_loss, epoch_acc
    
    def validate(self, val_loader: DataLoader, criterion) -> Tuple[float, float]:
        """Validate the model"""
        
        self.model.eval()
        running_loss = 0.0
        correct = 0
        total = 0
        
        with torch.no_grad():
            for images, labels in val_loader:
                images, labels = images.to(self.device), labels.to(self.device)
                
                outputs = self.model(images)
                loss = criterion(outputs, labels)
                
                running_loss += loss.item()
                _, predicted = torch.max(outputs.data, 1)
                total += labels.size(0)
                correct += (predicted == labels).sum().item()
        
        epoch_loss = running_loss / len(val_loader)
        epoch_acc = correct / total
        
        return epoch_loss, epoch_acc
    
    def train(self, data_dir: str, use_wandb: bool = False):
        """
        Main training loop
        
        Args:
            data_dir: Directory containing training data
            use_wandb: Whether to use Weights & Biases for logging
        """
        # Initialize wandb if requested
        if use_wandb:
            wandb.init(
                project="damaged-poster-detection",
                config=self.config.__dict__
            )
            wandb.watch(self.model)
        
        # Prepare data
        train_loader, val_loader, test_loader = self.prepare_data(data_dir)
        
        # Setup training components
        criterion = nn.CrossEntropyLoss()
        optimizer = optim.AdamW(
            self.model.parameters(),
            lr=self.config.learning_rate,
            weight_decay=self.config.weight_decay
        )
        
        # Learning rate scheduler
        scheduler = optim.lr_scheduler.ReduceLROnPlateau(
            optimizer, mode='min', factor=0.5, patience=5, verbose=True
        )
        
        # Early stopping
        best_val_loss = float('inf')
        patience_counter = 0
        
        logger.info("Starting training...")
        
        for epoch in range(self.config.num_epochs):
            logger.info(f"Epoch {epoch+1}/{self.config.num_epochs}")
            
            # Train
            train_loss, train_acc = self.train_epoch(train_loader, optimizer, criterion)
            
            # Validate
            val_loss, val_acc = self.validate(val_loader, criterion)
            
            # Update learning rate
            scheduler.step(val_loss)
            
            # Log metrics
            self.history['train_loss'].append(train_loss)
            self.history['train_acc'].append(train_acc)
            self.history['val_loss'].append(val_loss)
            self.history['val_acc'].append(val_acc)
            
            logger.info(f"Train Loss: {train_loss:.4f}, Train Acc: {train_acc:.4f}")
            logger.info(f"Val Loss: {val_loss:.4f}, Val Acc: {val_acc:.4f}")
            
            if use_wandb:
                wandb.log({
                    'train_loss': train_loss,
                    'train_acc': train_acc,
                    'val_loss': val_loss,
                    'val_acc': val_acc,
                    'learning_rate': optimizer.param_groups[0]['lr']
                })
            
            # Early stopping and model saving
            if val_loss < best_val_loss:
                best_val_loss = val_loss
                patience_counter = 0
                
                # Save best model
                model_path = os.path.join(
                    self.config.model_save_path, 
                    f"best_model_epoch_{epoch+1}.pth"
                )
                torch.save({
                    'epoch': epoch,
                    'model_state_dict': self.model.state_dict(),
                    'optimizer_state_dict': optimizer.state_dict(),
                    'val_loss': val_loss,
                    'config': self.config
                }, model_path)
                
                logger.info(f"Saved best model to {model_path}")
                
            else:
                patience_counter += 1
                
            if patience_counter >= self.config.patience:
                logger.info(f"Early stopping at epoch {epoch+1}")
                break
        
        # Final evaluation on test set
        test_loss, test_acc = self.validate(test_loader, criterion)
        logger.info(f"Final Test Loss: {test_loss:.4f}, Test Acc: {test_acc:.4f}")
        
        if use_wandb:
            wandb.log({'test_loss': test_loss, 'test_acc': test_acc})
            wandb.finish()
        
        # Save training history
        history_path = os.path.join(self.config.log_dir, "training_history.json")
        with open(history_path, 'w') as f:
            json.dump(self.history, f, indent=2)
        
        # Generate evaluation report
        self.evaluate_model(test_loader)
        
        logger.info("Training completed!")
    
    def evaluate_model(self, test_loader: DataLoader):
        """Generate detailed evaluation report"""
        
        self.model.eval()
        all_predictions = []
        all_labels = []
        
        with torch.no_grad():
            for images, labels in test_loader:
                images, labels = images.to(self.device), labels.to(self.device)
                outputs = self.model(images)
                _, predicted = torch.max(outputs, 1)
                
                all_predictions.extend(predicted.cpu().numpy())
                all_labels.extend(labels.cpu().numpy())
        
        # Classification report
        damage_types = [
            'no_damage', 'tear', 'crease', 'stain', 
            'fade', 'burn', 'water_damage'
        ]
        
        report = classification_report(
            all_labels, all_predictions, 
            target_names=damage_types,
            output_dict=True
        )
        
        # Save report
        report_path = os.path.join(self.config.log_dir, "evaluation_report.json")
        with open(report_path, 'w') as f:
            json.dump(report, f, indent=2)
        
        # Confusion matrix
        cm = confusion_matrix(all_labels, all_predictions)
        
        plt.figure(figsize=(10, 8))
        sns.heatmap(
            cm, annot=True, fmt='d', cmap='Blues',
            xticklabels=damage_types,
            yticklabels=damage_types
        )
        plt.title('Confusion Matrix')
        plt.xlabel('Predicted')
        plt.ylabel('Actual')
        plt.tight_layout()
        
        cm_path = os.path.join(self.config.log_dir, "confusion_matrix.png")
        plt.savefig(cm_path, dpi=300, bbox_inches='tight')
        plt.close()
        
        logger.info(f"Evaluation report saved to {report_path}")
        logger.info(f"Confusion matrix saved to {cm_path}")
    
    def plot_training_history(self):
        """Plot training history"""
        
        fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(15, 5))
        
        # Loss plot
        ax1.plot(self.history['train_loss'], label='Train Loss')
        ax1.plot(self.history['val_loss'], label='Val Loss')
        ax1.set_title('Model Loss')
        ax1.set_xlabel('Epoch')
        ax1.set_ylabel('Loss')
        ax1.legend()
        ax1.grid(True)
        
        # Accuracy plot
        ax2.plot(self.history['train_acc'], label='Train Accuracy')
        ax2.plot(self.history['val_acc'], label='Val Accuracy')
        ax2.set_title('Model Accuracy')
        ax2.set_xlabel('Epoch')
        ax2.set_ylabel('Accuracy')
        ax2.legend()
        ax2.grid(True)
        
        plt.tight_layout()
        
        plot_path = os.path.join(self.config.log_dir, "training_history.png")
        plt.savefig(plot_path, dpi=300, bbox_inches='tight')
        plt.close()
        
        logger.info(f"Training history plot saved to {plot_path}")


def create_sample_data_structure(base_dir: str):
    """
    Create sample directory structure for training data
    
    Args:
        base_dir: Base directory to create structure in
    """
    damage_types = [
        'no_damage', 'tear', 'crease', 'stain', 
        'fade', 'burn', 'water_damage'
    ]
    
    os.makedirs(base_dir, exist_ok=True)
    
    for damage_type in damage_types:
        damage_dir = os.path.join(base_dir, damage_type)
        os.makedirs(damage_dir, exist_ok=True)
    
    logger.info(f"Created sample data structure in {base_dir}")
    logger.info("Please organize your training images into the appropriate subdirectories:")
    for damage_type in damage_types:
        logger.info(f"  - {os.path.join(base_dir, damage_type)}/")


def main():
    """Main training function"""
    
    # Configuration
    config = TrainingConfig(
        data_dir="data/train",
        batch_size=16,  # Adjust based on GPU memory
        num_epochs=30,
        learning_rate=1e-4,
        use_augmentation=True
    )
    
    # Create sample data structure if needed
    if not os.path.exists(config.data_dir):
        create_sample_data_structure(config.data_dir)
        logger.info("Sample data structure created. Please add your training images and run again.")
        return
    
    # Initialize trainer
    trainer = DamageTrainer(config)
    
    # Start training
    trainer.train(config.data_dir, use_wandb=False)  # Set to True if you want to use wandb
    
    # Plot training history
    trainer.plot_training_history()


if __name__ == "__main__":
    main()