train.py

"""
Complete YOLOv8 Training, Evaluation and Prediction Script
Author: LenIAC
Date: 2024
Description: End-to-end implementation for YOLOv8 object detection
"""

import os
import yaml
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms
import numpy as np
from PIL import Image
import cv2
import random
from pathlib import Path
from typing import Dict, List, Tuple, Optional, Union
import matplotlib.pyplot as plt
from tqdm import tqdm
import warnings
warnings.filterwarnings('ignore')

from yolov8 import YOLOv8
from loss_fn import ComputeLoss

# Import the provided modules (assuming they're in the same directory)
# Note: We're not copying the model and loss function as requested
# They should be imported from the provided files

# ============================================================================
# DATA LOADING AND PREPROCESSING
# ============================================================================

class YOLODataset(Dataset):
    """
    Custom Dataset for YOLO format data.

    This class handles loading images and labels in YOLO format,
    applies augmentations, and prepares data for training.
    """

    def __init__(self,
                 data_yaml_path: str,
                 split: str = 'train',
                 img_size: int = 640,
                 augment: bool = True):
        """
        Initialize YOLO dataset.

        Args:
            data_yaml_path: Path to data.yaml file
            split: 'train' or 'val'
            img_size: Input image size (square)
            augment: Whether to apply data augmentation
        """
        super().__init__()

        self.img_size = img_size
        self.augment = augment and split == 'train'

        data_yaml_path = Path(data_yaml_path)
        self.datasets_dir = data_yaml_path.parent

        # Load data configuration
        with open(data_yaml_path, 'r') as f:
            data_config = yaml.safe_load(f)

        # Get image paths
        if split == 'train':
            image_dir = data_config['train']
        else:
            image_dir = data_config['val']

        self.image_dir = self.datasets_dir / image_dir
        self.label_dir = self.image_dir.parent / 'labels'

        # Get class names
        self.class_names = data_config['names']
        self.num_classes = data_config['nc']

        # Collect all image files
        self.image_files = list(self.image_dir.glob('*.jpg')) + \
                          list(self.image_dir.glob('*.jpeg')) + \
                          list(self.image_dir.glob('*.png'))

        if not self.image_files:
            raise ValueError(f"No images found in {self.image_dir}")

        print(f"Found {len(self.image_files)} images for {split} split")

        # Define augmentations
        self.transform = self._get_transforms()

    def _get_transforms(self):
        """Get image transformations based on augmentation flag."""
        if self.augment:
            return transforms.Compose([
                transforms.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2, hue=0.1),
                transforms.ToTensor(),
                transforms.RandomHorizontalFlip(p=0.5),
            ])
        else:
            return transforms.Compose([
                transforms.ToTensor(),
            ])

    def __len__(self) -> int:
        """Return number of images in dataset."""
        return len(self.image_files)

    def load_image(self, image_path: Path) -> torch.Tensor:
        """
        Load and preprocess image.

        Args:
            image_path: Path to image file

        Returns:
            Preprocessed image tensor
        """
        # Load image
        img = Image.open(image_path).convert('RGB')

        # Resize maintaining aspect ratio with padding
        img, ratio, pad = self.letterbox(img, self.img_size)

        # Apply transformations
        img = self.transform(img)

        return img, ratio, pad

    def load_labels(self, label_path: Path, ratio: Tuple, pad: Tuple) -> torch.Tensor:
        """
        Load and preprocess labels.

        Args:
            label_path: Path to label file
            ratio: Resize ratio (width_ratio, height_ratio)
            pad: Padding (width_pad, height_pad)

        Returns:
            Tensor of labels [num_objects, 5] where 5 = [class, x, y, w, h]
        """
        if not label_path.exists():
            return torch.zeros((0, 5))

        labels = []
        with open(label_path, 'r') as f:
            for line in f:
                parts = line.strip().split()
                if len(parts) == 5:
                    class_id = int(parts[0])
                    x_center = float(parts[1])
                    y_center = float(parts[2])
                    width = float(parts[3])
                    height = float(parts[4])

                    # Adjust for letterbox
                    x_center = x_center * ratio[0] + pad[0]
                    y_center = y_center * ratio[1] + pad[1]
                    width = width * ratio[0]
                    height = height * ratio[1]

                    # Normalize to image size
                    x_center /= self.img_size
                    y_center /= self.img_size
                    width /= self.img_size
                    height /= self.img_size

                    labels.append([class_id, x_center, y_center, width, height])

        if labels:
            return torch.tensor(labels)
        else:
            return torch.zeros((0, 5))

    def letterbox(self, img: Image.Image, new_shape: int = 640) -> Tuple:
        """
        Resize image with padding to maintain aspect ratio.

        Args:
            img: Input PIL Image
            new_shape: Target size

        Returns:
            Resized image, ratio, padding
        """
        # Current shape
        shape = img.size  # (width, height)

        # Calculate ratio
        r = min(new_shape / shape[0], new_shape / shape[1])

        # New dimensions
        new_unpad = (int(round(shape[0] * r)), int(round(shape[1] * r)))

        # Calculate padding
        dw = new_shape - new_unpad[0]
        dh = new_shape - new_unpad[1]
        dw /= 2
        dh /= 2

        # Resize
        img = img.resize(new_unpad, Image.BILINEAR)

        # Create new image with padding
        new_img = Image.new('RGB', (new_shape, new_shape), (114, 114, 114))
        new_img.paste(img, (int(round(dw)), int(round(dh))))

        return new_img, (r, r), (dw/new_shape, dh/new_shape)

    def __getitem__(self, idx: int) -> Dict:
        """
        Get single data sample.

        Args:
            idx: Index of sample

        Returns:
            Dictionary containing image and labels
        """
        # Get image path
        img_path = self.image_files[idx]

        # Load image
        image, ratio, pad = self.load_image(img_path)

        # Load corresponding labels
        label_path = self.label_dir / f"{img_path.stem}.txt"
        labels = self.load_labels(label_path, ratio, pad)

        return {
            'image': image,
            'labels': labels,
            'image_path': str(img_path)
        }

def collate_fn(batch: List[Dict]) -> Dict:
    """
    Custom collate function for DataLoader.

    Args:
        batch: List of samples

    Returns:
        Batched data
    """
    images = []
    labels = []
    image_paths = []

    for sample in batch:
        images.append(sample['image'])
        labels.append(sample['labels'])
        image_paths.append(sample['image_path'])

    # Stack images
    images = torch.stack(images, dim=0)

    return {
        'images': images,
        'labels': labels,
        'image_paths': image_paths
    }

# ============================================================================
# TRAINING UTILITIES
# ============================================================================

class ModelEMA:
    """
    Exponential Moving Average for model weights.

    This helps stabilize training and improve final model performance.
    """

    def __init__(self, model: nn.Module, decay: float = 0.9999):
        """
        Initialize EMA.

        Args:
            model: Model to apply EMA to
            decay: EMA decay rate
        """
        self.model = model
        self.decay = decay
        self.shadow = {}
        self.backup = {}

        # Register parameters
        for name, param in model.named_parameters():
            if param.requires_grad:
                self.shadow[name] = param.data.clone()

    def update(self):
        """Update EMA parameters."""
        for name, param in self.model.named_parameters():
            if param.requires_grad:
                assert name in self.shadow
                new_average = (1.0 - self.decay) * param.data + self.decay * self.shadow[name]
                self.shadow[name] = new_average.clone()

    def apply_shadow(self):
        """Apply EMA parameters to model."""
        for name, param in self.model.named_parameters():
            if param.requires_grad:
                assert name in self.shadow
                self.backup[name] = param.data.clone()
                param.data = self.shadow[name]

    def restore(self):
        """Restore original parameters."""
        for name, param in self.model.named_parameters():
            if param.requires_grad:
                assert name in self.backup
                param.data = self.backup[name]
        self.backup = {}

class EarlyStopping:
    """
    Early stopping to prevent overfitting.

    Stops training when validation loss doesn't improve for a given patience.
    """

    def __init__(self, patience: int = 10, min_delta: float = 0.0):
        """
        Initialize early stopping.

        Args:
            patience: Number of epochs to wait for improvement
            min_delta: Minimum change to qualify as improvement
        """
        self.patience = patience
        self.min_delta = min_delta
        self.counter = 0
        self.best_loss = None
        self.early_stop = False

    def __call__(self, val_loss: float) -> bool:
        """
        Check if training should stop.

        Args:
            val_loss: Current validation loss

        Returns:
            True if training should stop
        """
        if self.best_loss is None:
            self.best_loss = val_loss
            return False

        if val_loss > self.best_loss - self.min_delta:
            self.counter += 1
            if self.counter >= self.patience:
                self.early_stop = True
        else:
            self.best_loss = val_loss
            self.counter = 0

        return self.early_stop

# ============================================================================
# METRICS AND EVALUATION
# ============================================================================

class MetricsCalculator:
    """
    Calculate evaluation metrics for object detection.

    Computes mAP, precision, recall, and other relevant metrics.
    """

    def __init__(self, num_classes: int, iou_threshold: float = 0.5):
        """
        Initialize metrics calculator.

        Args:
            num_classes: Number of object classes
            iou_threshold: IoU threshold for considering detection as correct
        """
        self.num_classes = num_classes
        self.iou_threshold = iou_threshold
        self.reset()

    def reset(self):
        """Reset all metrics."""
        self.true_positives = []
        self.false_positives = []
        self.confidence_scores = []
        self.num_ground_truths = [0] * self.num_classes

    def update(self,
               predictions: List[torch.Tensor],
               ground_truths: List[torch.Tensor]):
        """
        Update metrics with new batch.

        Args:
            predictions: List of prediction tensors [N, 6] where 6 = [x1, y1, x2, y2, conf, class]
            ground_truths: List of ground truth tensors [M, 5] where 5 = [class, x, y, w, h]
        """
        for preds, gts in zip(predictions, ground_truths):
            if len(preds) == 0:
                continue

            if len(gts) == 0:
                # All predictions are false positives
                for pred in preds:
                    self.false_positives.append(pred[4].item())
                    self.confidence_scores.append(pred[4].item())
                continue

            # Convert ground truth format
            gt_boxes = []
            gt_classes = []
            for gt in gts:
                class_id = int(gt[0])
                x_center, y_center, width, height = gt[1:5].tolist()

                x1 = x_center - width / 2
                y1 = y_center - height / 2
                x2 = x_center + width / 2
                y2 = y_center + height / 2

                gt_boxes.append([x1, y1, x2, y2])
                gt_classes.append(class_id)
                self.num_ground_truths[class_id] += 1

            if not gt_boxes:
                continue

            gt_boxes = torch.tensor(gt_boxes)
            gt_classes = torch.tensor(gt_classes)

            # Sort predictions by confidence
            sorted_indices = torch.argsort(preds[:, 4], descending=True)
            preds = preds[sorted_indices]

            # Track which ground truths have been matched
            gt_matched = torch.zeros(len(gt_boxes), dtype=torch.bool)

            for pred in preds:
                pred_box = pred[:4].unsqueeze(0)
                pred_class = int(pred[5])
                confidence = pred[4].item()

                self.confidence_scores.append(confidence)

                # Find matching ground truth
                if len(gt_boxes) > 0:
                    # Get IoU with all ground truths of same class
                    same_class_mask = (gt_classes == pred_class)
                    if same_class_mask.any():
                        ious = self.calculate_iou(pred_box, gt_boxes[same_class_mask])
                        best_iou, best_idx = torch.max(ious, dim=1)

                        if best_iou.item() >= self.iou_threshold:
                            # Get original index
                            original_idx = torch.where(same_class_mask)[0][best_idx.item()]

                            if not gt_matched[original_idx]:
                                # True positive
                                self.true_positives.append(confidence)
                                gt_matched[original_idx] = True
                                continue

                # False positive
                self.false_positives.append(confidence)

    def calculate_iou(self, box1: torch.Tensor, box2: torch.Tensor) -> torch.Tensor:
        """
        Calculate IoU between boxes.

        Args:
            box1: Tensor of shape [1, 4]
            box2: Tensor of shape [N, 4]

        Returns:
            IoU values
        """
        # Get coordinates
        x1 = torch.max(box1[:, 0], box2[:, 0])
        y1 = torch.max(box1[:, 1], box2[:, 1])
        x2 = torch.min(box1[:, 2], box2[:, 2])
        y2 = torch.min(box1[:, 3], box2[:, 3])

        # Intersection area
        intersection = torch.clamp(x2 - x1, min=0) * torch.clamp(y2 - y1, min=0)

        # Union area
        area1 = (box1[:, 2] - box1[:, 0]) * (box1[:, 3] - box1[:, 1])
        area2 = (box2[:, 2] - box2[:, 0]) * (box2[:, 3] - box2[:, 1])
        union = area1 + area2 - intersection

        # IoU
        iou = intersection / union
        return iou

    def compute_map(self) -> float:
        """
        Compute mean Average Precision.

        Returns:
            mAP value
        """
        if not self.confidence_scores:
            return 0.0

        # Sort by confidence
        sorted_indices = np.argsort(self.confidence_scores)[::-1]

        # Convert to numpy arrays
        tp = np.array([1 if score in self.true_positives else 0
                      for score in self.confidence_scores])
        fp = np.array([1 if score in self.false_positives else 0
                      for score in self.confidence_scores])

        # Sort by confidence
        tp = tp[sorted_indices]
        fp = fp[sorted_indices]

        # Compute precision-recall curve
        tp_cumsum = np.cumsum(tp)
        fp_cumsum = np.cumsum(fp)

        recalls = tp_cumsum / (sum(self.num_ground_truths) + 1e-16)
        precisions = tp_cumsum / (tp_cumsum + fp_cumsum + 1e-16)

        # Compute AP using 101-point interpolation
        ap = 0.0
        for t in np.arange(0, 1.01, 0.01):
            mask = recalls >= t
            if mask.any():
                ap += np.max(precisions[mask])

        return ap / 101

# ============================================================================
# TRAINING LOOP
# ============================================================================

class YOLOv8Trainer:
    """
    Complete trainer for YOLOv8 model.

    Handles training, validation, and model saving.
    """

    def __init__(self,
                 model: nn.Module,
                 train_loader: DataLoader,
                 val_loader: DataLoader,
                 config: Dict):
        """
        Initialize trainer.

        Args:
            model: YOLOv8 model
            train_loader: Training data loader
            val_loader: Validation data loader
            config: Training configuration
        """
        self.model = model
        self.train_loader = train_loader
        self.val_loader = val_loader
        self.config = config

        # Setup device
        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        self.model.to(self.device)

        self.loss_fn = ComputeLoss(self.model, self.config)

        # Setup optimizer
        self.optimizer = optim.AdamW(
            model.parameters(), lr=config['learning_rate'], weight_decay=config['weight_decay']
        )

        # Setup learning rate scheduler
        self.scheduler = optim.lr_scheduler.CosineAnnealingLR(
            self.optimizer,
            T_max=config['epochs'],
            eta_min=config['learning_rate'] * 0.01
        )

        # Setup EMA
        self.ema = ModelEMA(model, decay=config['ema_decay'])

        # Setup early stopping
        self.early_stopping = EarlyStopping(
            patience=config['early_stopping_patience'],
            min_delta=config['early_stopping_min_delta']
        )

        # Setup metrics
        self.metrics_calculator = MetricsCalculator(
            num_classes=config['num_classes']
        )

        # Training state
        self.current_epoch = 0
        self.best_map = 0.0
        self.train_losses = []
        self.val_losses = []
        self.val_maps = []

        # Create output directory
        self.output_dir = Path(config['output_dir'])
        self.output_dir.mkdir(parents=True, exist_ok=True)

        # Save configuration
        with open(self.output_dir / 'config.yaml', 'w') as f:
            yaml.dump(config, f)

    def train_epoch(self) -> float:
        """
        Train for one epoch.

        Returns:
            Average training loss
        """
        self.model.train()
        total_loss = 0.0
        num_batches = 0

        progress_bar = tqdm(self.train_loader, desc=f"Epoch {self.current_epoch + 1}")

        for batch_idx, batch in enumerate(progress_bar):
            # Move data to device
            images = batch['images'].to(self.device)
            labels = batch['labels']

            # Prepare targets
            targets = self.prepare_targets(labels)

            # Forward pass
            self.optimizer.zero_grad()
            outputs = self.model(images)

            # Compute loss
            # Note: We assume the loss function is available
            # In practice, you would import it from loss_fn.py
            progress_bar.write("outputs: " + str(outputs) + " targets: " + str(targets))
            loss_box, loss_cls, loss_dfl = self.compute_loss(outputs, targets)

            total_loss = loss_box + loss_cls + loss_dfl

            # Backward pass
            total_loss.backward()

            # Gradient clipping
            torch.nn.utils.clip_grad_norm_(self.model.parameters(),max_norm=self.config['grad_clip'])

            # Optimizer step
            self.optimizer.step()

            # Update EMA
            self.ema.update()

            # Update progress bar
            progress_bar.set_postfix({
                'loss': total_loss.item(),
                'box': loss_box.item(),
                'cls': loss_cls.item(),
                'dfl': loss_dfl.item()
            })

            total_loss += total_loss.item()
            num_batches += 1

        return total_loss / num_batches

    def validate(self) -> Tuple[float, float]:
        """
        Validate model.

        Returns:
            Tuple of (validation loss, mAP)
        """
        self.model.eval()
        self.ema.apply_shadow()

        total_loss = 0.0
        num_batches = 0
        all_predictions = []
        all_ground_truths = []

        with torch.no_grad():
            progress_bar = tqdm(self.val_loader, desc="Validation")

            for batch_idx, batch in enumerate(progress_bar):
                # Move data to device
                images = batch['images'].to(self.device)
                labels = batch['labels']

                # Prepare targets
                targets = self.prepare_targets(labels)

                # Forward pass
                outputs = self.model(images)

                # Compute loss
                progress_bar.write("outputs: " + str(outputs) + " targets: " + str(targets))
                loss_box, loss_cls, loss_dfl = self.compute_loss(outputs, targets)
                batch_loss = loss_box + loss_cls + loss_dfl

                total_loss += batch_loss.item()
                num_batches += 1

                # Get predictions for metrics
                predictions = self.get_predictions(outputs)
                all_predictions.extend(predictions)
                all_ground_truths.extend(labels)

                progress_bar.set_postfix({
                    'val_loss': batch_loss.item()
                })

        # Restore original model
        self.ema.restore()

        # Calculate mAP
        self.metrics_calculator.reset()
        self.metrics_calculator.update(all_predictions, all_ground_truths)
        mAP = self.metrics_calculator.compute_map()

        return total_loss / num_batches, mAP

    def prepare_targets(self, labels: List[torch.Tensor]) -> Dict:
        """
        Prepare targets for loss computation.

        Args:
            labels: List of label tensors

        Returns:
            Dictionary of targets
        """
        # This is a simplified version
        # In practice, you would implement the full target preparation
        # as in the loss_fn.py file

        batch_size = len(labels)
        max_objects = max(len(l) for l in labels) if labels else 0

        if max_objects == 0:
            return {
                'idx': torch.zeros((0, 1)),
                'cls': torch.zeros((0, 1)),
                'box': torch.zeros((0, 4))
            }

        # Initialize tensors
        idx_tensor = torch.zeros((batch_size * max_objects, 1))
        cls_tensor = torch.zeros((batch_size * max_objects, 1))
        box_tensor = torch.zeros((batch_size * max_objects, 4))

        idx = 0
        for batch_idx, label_batch in enumerate(labels):
            for obj_idx in range(len(label_batch)):
                label = label_batch[obj_idx]
                idx_tensor[idx] = batch_idx
                cls_tensor[idx] = label[0]

                # Convert from center format to corner format
                x_center, y_center, width, height = label[1:5]
                x1 = x_center - width / 2
                y1 = y_center - height / 2
                x2 = x_center + width / 2
                y2 = y_center + height / 2

                box_tensor[idx] = torch.tensor([x1, y1, x2, y2])
                idx += 1

        # Trim to actual number of objects
        idx_tensor = idx_tensor[:idx]
        cls_tensor = cls_tensor[:idx]
        box_tensor = box_tensor[:idx]

        return {
            'idx': idx_tensor,
            'cls': cls_tensor,
            'box': box_tensor
        }

    def compute_loss(self, outputs, targets):
        """
        Compute loss using the provided loss function.

        Args:
            outputs: Model outputs
            targets: Prepared targets

        Returns:
            Tuple of (box_loss, cls_loss, dfl_loss)
        """
        # Note: In practice, you would import and use the ComputeLoss class
        # from loss_fn.py. This is a simplified placeholder.

        # For demonstration, we'll use placeholder losses
        box_loss, cls_loss, dfl_loss = self.loss_fn(outputs, targets)

        # Placeholder implementation
        # device = outputs[0].device
        # box_loss = torch.tensor(0.5, device=device, requires_grad=True)
        # cls_loss = torch.tensor(0.3, device=device, requires_grad=True)
        # dfl_loss = torch.tensor(0.2, device=device, requires_grad=True)

        return box_loss, cls_loss, dfl_loss

    def get_predictions(self, outputs, confidence_threshold: float = 0.25):
        """
        Convert model outputs to predictions.

        Args:
            outputs: Model outputs
            confidence_threshold: Minimum confidence score

        Returns:
            List of predictions per image
        """
        predictions = []

        # Note: This is a simplified version
        # In practice, use the model's head for inference

        for output in outputs:
            # output shape: [batch_size, num_predictions, 85] for COCO
            # where 85 = 4 box coordinates + 1 objectness + 80 classes

            batch_predictions = []
            for img_idx in range(output.shape[0]):
                img_preds = output[img_idx]

                # Filter by confidence
                if img_preds.shape[1] > 5:  # Has class predictions
                    # Get objectness and class scores
                    obj_scores = img_preds[:, 4:5]
                    cls_scores = img_preds[:, 5:]

                    # Combine scores
                    scores = obj_scores * cls_scores.max(dim=1, keepdim=True)[0]

                    # Filter by threshold
                    mask = scores.squeeze() > confidence_threshold
                    filtered_preds = img_preds[mask]

                    if len(filtered_preds) > 0:
                        # Get boxes and classes
                        boxes = filtered_preds[:, :4]
                        confidences = scores[mask]
                        class_ids = filtered_preds[:, 5:].argmax(dim=1)

                        # Combine into [x1, y1, x2, y2, conf, class]
                        batch_preds = torch.cat([boxes, confidences, class_ids.unsqueeze(1).float()], dim=1)
                        batch_predictions.append(batch_preds)
                    else:
                        batch_predictions.append(torch.zeros((0, 6)))
                else:
                    batch_predictions.append(torch.zeros((0, 6)))

            predictions.extend(batch_predictions)

        return predictions

    def save_checkpoint(self, filename: str, is_best: bool = False):
        """
        Save model checkpoint.

        Args:
            filename: Checkpoint filename
            is_best: Whether this is the best model
        """
        checkpoint = {
            'epoch': self.current_epoch,
            'model_state_dict': self.model.state_dict(),
            'optimizer_state_dict': self.optimizer.state_dict(),
            'scheduler_state_dict': self.scheduler.state_dict(),
            'best_map': self.best_map,
            'train_losses': self.train_losses,
            'val_losses': self.val_losses,
            'val_maps': self.val_maps,
            'config': self.config
        }

        torch.save(checkpoint, self.output_dir / filename)

        if is_best:
            torch.save(self.model.state_dict(), self.output_dir / 'best_model.pth')

    def load_checkpoint(self, checkpoint_path: str):
        """
        Load model checkpoint.

        Args:
            checkpoint_path: Path to checkpoint file
        """
        checkpoint = torch.load(checkpoint_path, map_location=self.device)

        self.model.load_state_dict(checkpoint['model_state_dict'])
        self.optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
        self.scheduler.load_state_dict(checkpoint['scheduler_state_dict'])

        self.current_epoch = checkpoint['epoch']
        self.best_map = checkpoint['best_map']
        self.train_losses = checkpoint['train_losses']
        self.val_losses = checkpoint['val_losses']
        self.val_maps = checkpoint['val_maps']

        print(f"Loaded checkpoint from epoch {self.current_epoch}")
        print(f"Best mAP: {self.best_map:.4f}")

    def plot_training_history(self):
        """Plot training history."""
        fig, axes = plt.subplots(1, 2, figsize=(12, 4))

        # Plot losses
        axes[0].plot(self.train_losses, label='Training Loss', color='blue')
        axes[0].plot(self.val_losses, label='Validation Loss', color='red')
        axes[0].set_xlabel('Epoch')
        axes[0].set_ylabel('Loss')
        axes[0].set_title('Training and Validation Loss')
        axes[0].legend()
        axes[0].grid(True)

        # Plot mAP
        axes[1].plot(self.val_maps, label='Validation mAP', color='green')
        axes[1].set_xlabel('Epoch')
        axes[1].set_ylabel('mAP')
        axes[1].set_title('Validation mAP')
        axes[1].legend()
        axes[1].grid(True)

        plt.tight_layout()
        plt.savefig(self.output_dir / 'training_history.png', dpi=150)
        plt.close()

    def train(self):
        """Main training loop."""
        print(f"Starting training on {self.device}")
        print(f"Number of parameters: {sum(p.numel() for p in self.model.parameters()):,}")

        for epoch in range(self.current_epoch, self.config['epochs']):
            self.current_epoch = epoch

            # Train for one epoch
            train_loss = self.train_epoch()
            self.train_losses.append(train_loss)

            # Validate
            val_loss, val_map = self.validate()
            self.val_losses.append(val_loss)
            self.val_maps.append(val_map)

            # Update learning rate
            self.scheduler.step()

            # Print epoch summary
            print(f"\nEpoch {epoch + 1}/{self.config['epochs']}:")
            print(f"  Train Loss: {train_loss:.4f}")
            print(f"  Val Loss: {val_loss:.4f}")
            print(f"  Val mAP: {val_map:.4f}")
            print(f"  Learning Rate: {self.scheduler.get_last_lr()[0]:.6f}")

            # Save checkpoint
            self.save_checkpoint(f'checkpoint_epoch_{epoch + 1}.pth')

            # Save best model
            if val_map > self.best_map:
                self.best_map = val_map
                self.save_checkpoint('best_checkpoint.pth', is_best=True)
                print(f"  New best model saved with mAP: {val_map:.4f}")

            # Check early stopping
            if self.early_stopping(val_loss):
                print(f"Early stopping triggered at epoch {epoch + 1}")
                break

            # Plot training history
            if (epoch + 1) % 5 == 0:
                self.plot_training_history()

        # Final plot
        self.plot_training_history()
        print(f"\nTraining completed. Best mAP: {self.best_map:.4f}")

# ============================================================================
# PREDICTION AND INFERENCE
# ============================================================================

class YOLOv8Predictor:
    """
    Predictor for YOLOv8 model.

    Handles loading trained models and making predictions on new images.
    """

    def __init__(self,
                 model_path: str,
                 config_path: str,
                 confidence_threshold: float = 0.25,
                 iou_threshold: float = 0.45):
        """
        Initialize predictor.

        Args:
            model_path: Path to trained model weights
            config_path: Path to config file
            confidence_threshold: Minimum confidence for predictions
            iou_threshold: IoU threshold for NMS
        """
        # Load configuration
        with open(config_path, 'r') as f:
            self.config = yaml.safe_load(f)

        # Initialize model
        # Note: You need to import your YOLOv8 model class
        # from yolov8 import YOLOv8
        self.model = self._initialize_model()

        # Load weights
        self.load_weights(model_path)

        # Set thresholds
        self.confidence_threshold = confidence_threshold
        self.iou_threshold = iou_threshold

        # Setup device
        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        self.model.to(self.device)
        self.model.eval()

    def _initialize_model(self):
        """
        Initialize YOLOv8 model.

        Returns:
            Initialized model
        """
        # This is a placeholder - you need to import your actual model
        # Example: return YOLOv8(version='n', num_classes=self.config['num_classes'])

        # For demonstration, create a dummy model
        class DummyModel(nn.Module):
            def __init__(self):
                super().__init__()
                self.dummy = nn.Parameter(torch.randn(1))

            def forward(self, x):
                return [torch.randn(x.shape[0], 84, 8400)]  # Dummy output

        return DummyModel()

    def load_weights(self, model_path: str):
        """
        Load model weights.

        Args:
            model_path: Path to model weights
        """
        if os.path.exists(model_path):
            state_dict = torch.load(model_path, map_location='cpu')

            # Handle different checkpoint formats
            if 'model_state_dict' in state_dict:
                state_dict = state_dict['model_state_dict']

            # Load weights
            self.model.load_state_dict(state_dict, strict=False)
            print(f"Loaded weights from {model_path}")
        else:
            print(f"Warning: Model weights not found at {model_path}")

    def preprocess_image(self, image_path: Union[str, Path]) -> torch.Tensor:
        """
        Preprocess image for inference.

        Args:
            image_path: Path to image file

        Returns:
            Preprocessed image tensor
        """
        # Load image
        if isinstance(image_path, (str, Path)):
            img = Image.open(image_path).convert('RGB')
        else:
            # Assume it's already a PIL Image
            img = image_path

        # Get original dimensions
        original_size = img.size  # (width, height)

        # Resize with letterbox
        img_size = self.config.get('img_size', 640)
        img, ratio, pad = self.letterbox(img, img_size)

        # Convert to tensor
        transform = transforms.Compose([
            transforms.ToTensor(),
        ])
        img_tensor = transform(img).unsqueeze(0)  # Add batch dimension

        return img_tensor, original_size, ratio, pad

    def letterbox(self, img: Image.Image, new_shape: int = 640) -> Tuple:
        """
        Resize image with padding (same as in dataset).

        Args:
            img: Input PIL Image
            new_shape: Target size

        Returns:
            Resized image, ratio, padding
        """
        # Current shape
        shape = img.size  # (width, height)

        # Calculate ratio
        r = min(new_shape / shape[0], new_shape / shape[1])

        # New dimensions
        new_unpad = (int(round(shape[0] * r)), int(round(shape[1] * r)))

        # Calculate padding
        dw = new_shape - new_unpad[0]
        dh = new_shape - new_unpad[1]
        dw /= 2
        dh /= 2

        # Resize
        img = img.resize(new_unpad, Image.BILINEAR)

        # Create new image with padding
        new_img = Image.new('RGB', (new_shape, new_shape), (114, 114, 114))
        new_img.paste(img, (int(round(dw)), int(round(dh))))

        return new_img, (r, r), (dw/new_shape, dh/new_shape)

    def postprocess_predictions(self,
                                predictions: torch.Tensor,
                                original_size: Tuple[int, int],
                                ratio: Tuple[float, float],
                                pad: Tuple[float, float]) -> List[Dict]:
        """
        Postprocess model predictions.

        Args:
            predictions: Raw model predictions
            original_size: Original image size (width, height)
            ratio: Resize ratio
            pad: Padding applied

        Returns:
            List of detected objects with bounding boxes and confidence scores
        """
        if predictions.shape[1] == 0:
            return []

        # Apply confidence threshold
        if predictions.shape[2] > 5:  # Has class predictions
            # Get objectness and class scores
            obj_scores = predictions[..., 4:5]
            cls_scores = predictions[..., 5:]

            # Combine scores
            scores = obj_scores * cls_scores.max(dim=-1, keepdim=True)[0]

            # Filter by threshold
            mask = scores.squeeze() > self.confidence_threshold
            filtered_preds = predictions[:, mask, :]
        else:
            filtered_preds = predictions

        if filtered_preds.shape[1] == 0:
            return []

        # Apply NMS
        boxes = filtered_preds[..., :4]
        scores = filtered_preds[..., 4]

        # Convert to corner format if needed
        # Assuming boxes are in [x_center, y_center, width, height] format
        boxes_corners = torch.zeros_like(boxes)
        boxes_corners[..., 0] = boxes[..., 0] - boxes[..., 2] / 2  # x1
        boxes_corners[..., 1] = boxes[..., 1] - boxes[..., 3] / 2  # y1
        boxes_corners[..., 2] = boxes[..., 0] + boxes[..., 2] / 2  # x2
        boxes_corners[..., 3] = boxes[..., 1] + boxes[..., 3] / 2  # y2

        # Apply NMS
        keep_indices = self.non_max_suppression(boxes_corners[0], scores[0], self.iou_threshold)

        # Get final predictions
        final_predictions = filtered_preds[0, keep_indices]

        # Convert to original image coordinates
        results = []
        for pred in final_predictions:
            # Get box coordinates
            x_center, y_center, width, height = pred[:4].cpu().numpy()

            # Remove padding and rescale
            x_center = (x_center - pad[0]) / ratio[0]
            y_center = (y_center - pad[1]) / ratio[1]
            width = width / ratio[0]
            height = height / ratio[1]

            # Convert to corner coordinates
            x1 = max(0, x_center - width / 2)
            y1 = max(0, y_center - height / 2)
            x2 = min(original_size[0], x_center + width / 2)
            y2 = min(original_size[1], y_center + height / 2)

            # Get class and confidence
            if pred.shape[0] > 5:
                class_scores = pred[5:].cpu().numpy()
                class_id = np.argmax(class_scores)
                confidence = pred[4].item() * class_scores[class_id]
            else:
                class_id = 0
                confidence = pred[4].item()

            results.append({
                'bbox': [x1, y1, x2, y2],
                'confidence': float(confidence),
                'class_id': int(class_id),
                'class_name': self.config['class_names'][class_id]
                if 'class_names' in self.config else f'class_{class_id}'
            })

        return results

    def non_max_suppression(self, boxes: torch.Tensor, scores: torch.Tensor, iou_threshold: float) -> List[int]:
        """
        Apply Non-Maximum Suppression.

        Args:
            boxes: Bounding boxes [N, 4] in corner format
            scores: Confidence scores [N]
            iou_threshold: IoU threshold

        Returns:
            Indices of boxes to keep
        """
        if boxes.shape[0] == 0:
            return []

        # Sort by score
        sorted_indices = torch.argsort(scores, descending=True)
        boxes = boxes[sorted_indices]
        scores = scores[sorted_indices]

        keep = []

        while boxes.shape[0] > 0:
            # Keep the box with highest score
            keep.append(sorted_indices[0].item())

            if boxes.shape[0] == 1:
                break

            # Compute IoU with remaining boxes
            ious = self.compute_iou(boxes[0:1], boxes[1:])

            # Remove boxes with IoU > threshold
            mask = ious.squeeze() <= iou_threshold
            boxes = boxes[1:][mask]
            scores = scores[1:][mask]
            sorted_indices = sorted_indices[1:][mask]

        return keep

    def compute_iou(self, box1: torch.Tensor, box2: torch.Tensor) -> torch.Tensor:
        """
        Compute IoU between boxes.

        Args:
            box1: Tensor of shape [1, 4]
            box2: Tensor of shape [N, 4]

        Returns:
            IoU values
        """
        # Get coordinates
        x1 = torch.max(box1[:, 0], box2[:, 0])
        y1 = torch.max(box1[:, 1], box2[:, 1])
        x2 = torch.min(box1[:, 2], box2[:, 2])
        y2 = torch.min(box1[:, 3], box2[:, 3])

        # Intersection area
        intersection = torch.clamp(x2 - x1, min=0) * torch.clamp(y2 - y1, min=0)

        # Union area
        area1 = (box1[:, 2] - box1[:, 0]) * (box1[:, 3] - box1[:, 1])
        area2 = (box2[:, 2] - box2[:, 0]) * (box2[:, 3] - box2[:, 1])
        union = area1 + area2 - intersection

        # IoU
        iou = intersection / union
        return iou

    def predict(self, image_path: Union[str, Path]) -> Tuple[np.ndarray, List[Dict]]:
        """
        Make predictions on an image.

        Args:
            image_path: Path to image file

        Returns:
            Tuple of (image with drawn boxes, list of detections)
        """
        # Preprocess image
        img_tensor, original_size, ratio, pad = self.preprocess_image(image_path)
        img_tensor = img_tensor.to(self.device)

        # Make prediction
        with torch.no_grad():
            outputs = self.model(img_tensor)

        # Postprocess predictions
        detections = self.postprocess_predictions(outputs[0], original_size, ratio, pad)

        # Load original image for visualization
        if isinstance(image_path, (str, Path)):
            img = cv2.imread(str(image_path))
            img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        else:
            img = np.array(image_path)

        # Draw bounding boxes
        img_with_boxes = self.draw_boxes(img.copy(), detections)

        return img_with_boxes, detections

    def draw_boxes(self, image: np.ndarray, detections: List[Dict]) -> np.ndarray:
        """
        Draw bounding boxes on image.

        Args:
            image: Input image
            detections: List of detections

        Returns:
            Image with drawn boxes
        """
        img_with_boxes = image.copy()

        # Define colors for different classes
        colors = [
            (255, 0, 0),    # Red
            (0, 255, 0),    # Green
            (0, 0, 255),    # Blue
            (255, 255, 0),  # Cyan
            (255, 0, 255),  # Magenta
            (0, 255, 255),  # Yellow
        ]

        for det in detections:
            bbox = det['bbox']
            class_id = det['class_id']
            confidence = det['confidence']
            class_name = det['class_name']

            # Get color for this class
            color = colors[class_id % len(colors)]

            # Convert coordinates to integers
            x1, y1, x2, y2 = map(int, bbox)

            # Draw rectangle
            cv2.rectangle(img_with_boxes, (x1, y1), (x2, y2), color, 2)

            # Create label
            label = f"{class_name}: {confidence:.2f}"

            # Get text size
            (text_width, text_height), baseline = cv2.getTextSize(
                label, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 1
            )

            # Draw background for text
            cv2.rectangle(
                img_with_boxes,
                (x1, y1 - text_height - baseline - 5),
                (x1 + text_width, y1),
                color,
                -1
            )

            # Draw text
            cv2.putText(
                img_with_boxes,
                label,
                (x1, y1 - baseline - 5),
                cv2.FONT_HERSHEY_SIMPLEX,
                0.5,
                (255, 255, 255),
                1
            )

        return img_with_boxes

# ============================================================================
# MAIN TRAINING SCRIPT
# ============================================================================

def main():
    """Main training script."""

    # Configuration
    config = {
        # Model configuration
        'model_version': 'n',  # n, s, m, l, x
        'num_classes': 80,     # COCO has 80 classes

        # Training configuration
        'epochs': 100,
        'batch_size': 16,
        'learning_rate': 0.001,
        'weight_decay': 0.0005,
        'grad_clip': 10.0,

        # Data configuration
        'img_size': 640,
        'augment': True,

        # Loss weights
        'box_loss_weight': 7.5,
        'cls_loss_weight': 0.5,
        'dfl_loss_weight': 1.5,

        # EMA configuration
        'ema_decay': 0.9999,

        # Early stopping
        'early_stopping_patience': 20,
        'early_stopping_min_delta': 0.001,

        # Output configuration
        'output_dir': 'runs/train',
        'checkpoint_interval': 5,

        # Dataset paths
        'data_yaml': 'datasets/data.yaml',
    }

    # Set random seeds for reproducibility
    torch.manual_seed(42)
    np.random.seed(42)
    random.seed(42)

    if torch.cuda.is_available():
        torch.cuda.manual_seed(42)
        torch.cuda.manual_seed_all(42)
        torch.backends.cudnn.deterministic = True
        torch.backends.cudnn.benchmark = False

    # Create datasets
    print("Creating datasets...")
    train_dataset = YOLODataset(
        data_yaml_path=config['data_yaml'],
        split='train',
        img_size=config['img_size'],
        augment=config['augment']
    )

    val_dataset = YOLODataset(
        data_yaml_path=config['data_yaml'],
        split='val',
        img_size=config['img_size'],
        augment=False
    )

    # Create data loaders
    train_loader = DataLoader(
        train_dataset,
        batch_size=config['batch_size'],
        shuffle=True,
        num_workers=4,
        pin_memory=True,
        collate_fn=collate_fn,
        drop_last=True
    )

    val_loader = DataLoader(
        val_dataset,
        batch_size=config['batch_size'],
        shuffle=False,
        num_workers=4,
        pin_memory=True,
        collate_fn=collate_fn
    )

    # Initialize model
    print("Initializing model...")
    # Note: You need to import your YOLOv8 model
    # from yolov8 import YOLOv8
    model = YOLOv8(version=config['model_version'], num_classes=train_dataset.num_classes)

    # For demonstration, create a dummy model
    # model = nn.Sequential(
    #     nn.Conv2d(3, 64, kernel_size=3, padding=1),
    #     nn.ReLU(),
    #     nn.AdaptiveAvgPool2d(1)
    # )

    # Create trainer
    trainer = YOLOv8Trainer(model=model, train_loader=train_loader, val_loader=val_loader, config=config)

    # Start training
    trainer.train()

    print("\nTraining completed successfully!")
    print(f"Best mAP achieved: {trainer.best_map:.4f}")
    print(f"Model saved in: {trainer.output_dir}")

# ============================================================================
# EVALUATION SCRIPT
# ============================================================================

def evaluate_model(model_path: str, config_path: str, data_yaml: str):
    """
    Evaluate trained model.

    Args:
        model_path: Path to trained model
        config_path: Path to config file
        data_yaml: Path to data.yaml
    """
    # Load configuration
    with open(config_path, 'r') as f:
        config = yaml.safe_load(f)

    # Create validation dataset
    val_dataset = YOLODataset(
        data_yaml_path=data_yaml,
        split='val',
        img_size=config['img_size'],
        augment=False
    )

    val_loader = DataLoader(
        val_dataset,
        batch_size=config['batch_size'],
        shuffle=False,
        num_workers=4,
        pin_memory=True,
        collate_fn=collate_fn
    )

    # Initialize model
    # Note: Load your actual model here
    model = nn.Sequential(
        nn.Conv2d(3, 64, kernel_size=3, padding=1),
        nn.ReLU(),
        nn.AdaptiveAvgPool2d(1)
    )

    # Load weights
    checkpoint = torch.load(model_path, map_location='cpu')
    if 'model_state_dict' in checkpoint:
        model.load_state_dict(checkpoint['model_state_dict'])
    else:
        model.load_state_dict(checkpoint)

    # Move to device
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    model.to(device)
    model.eval()

    # Initialize metrics
    metrics_calculator = MetricsCalculator(num_classes=config['num_classes'])

    # Evaluate
    all_predictions = []
    all_ground_truths = []

    with torch.no_grad():
        progress_bar = tqdm(val_loader, desc="Evaluation")

        for batch in progress_bar:
            images = batch['images'].to(device)
            labels = batch['labels']

            # Forward pass
            outputs = model(images)

            # Get predictions (simplified)
            predictions = []
            for i in range(len(outputs)):
                # Dummy predictions for demonstration
                pred = torch.randn(10, 6)  # [num_preds, 6]
                pred[:, 4] = torch.sigmoid(torch.randn(10))  # Confidence
                pred[:, 5] = torch.randint(0, config['num_classes'], (10,))
                predictions.append(pred)

            all_predictions.extend(predictions)
            all_ground_truths.extend(labels)

    # Calculate metrics
    metrics_calculator.update(all_predictions, all_ground_truths)
    mAP = metrics_calculator.compute_map()

    print(f"\nEvaluation Results:")
    print(f"  mAP@0.5: {mAP:.4f}")
    print(f"  Number of validation images: {len(val_dataset)}")

    return mAP

# ============================================================================
# PREDICTION DEMONSTRATION
# ============================================================================

def demonstrate_predictions(model_path: str, config_path: str, image_dir: str):
    """
    Demonstrate predictions on sample images.

    Args:
        model_path: Path to trained model
        config_path: Path to config file
        image_dir: Directory with test images
    """
    # Initialize predictor
    predictor = YOLOv8Predictor(
        model_path=model_path,
        config_path=config_path,
        confidence_threshold=0.25,
        iou_threshold=0.45
    )

    # Get test images
    image_dir = Path(image_dir)
    test_images = list(image_dir.glob('*.jpg')) + \
                  list(image_dir.glob('*.jpeg')) + \
                  list(image_dir.glob('*.png'))

    if not test_images:
        print(f"No images found in {image_dir}")
        return

    # Make predictions on first few images
    num_demo = min(5, len(test_images))

    for i in range(num_demo):
        image_path = test_images[i]
        print(f"\nProcessing image {i+1}/{num_demo}: {image_path.name}")

        # Make prediction
        img_with_boxes, detections = predictor.predict(image_path)

        # Print results
        print(f"  Detected {len(detections)} objects:")
        for det in detections:
            print(f"    - {det['class_name']}: {det['confidence']:.2f} "
                  f"at [{det['bbox'][0]:.0f}, {det['bbox'][1]:.0f}, "
                  f"{det['bbox'][2]:.0f}, {det['bbox'][3]:.0f}]")

        # Save visualization
        output_path = Path('predictions') / f"pred_{image_path.name}"
        output_path.parent.mkdir(exist_ok=True)

        # Convert back to BGR for OpenCV save
        img_bgr = cv2.cvtColor(img_with_boxes, cv2.COLOR_RGB2BGR)
        cv2.imwrite(str(output_path), img_bgr)

        print(f"  Visualization saved to: {output_path}")

# ============================================================================
# ENTRY POINT
# ============================================================================

if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser(description="YOLOv8 Training and Evaluation")
    parser.add_argument('--mode', type=str, default='train',
                       choices=['train', 'eval', 'predict'],
                       help='Mode: train, eval, or predict')
    parser.add_argument('--model', type=str, default='runs/train/best_model.pth',
                       help='Path to model weights')
    parser.add_argument('--config', type=str, default='runs/train/config.yaml',
                       help='Path to config file')
    parser.add_argument('--data', type=str, default='datasets/data.yaml',
                       help='Path to data.yaml')
    parser.add_argument('--images', type=str, default='test_images',
                       help='Directory with test images for prediction mode')

    args = parser.parse_args()

    if args.mode == 'train':
        main()
    elif args.mode == 'eval':
        evaluate_model(args.model, args.config, args.data)
    elif args.mode == 'predict':
        demonstrate_predictions(args.model, args.config, args.images)
    else:
        print(f"Unknown mode: {args.mode}")