src/segmentation/trainer.py

"""

trainer.py

----------

YOLOv8 segmentation trainer for floor plan element detection.



Automatically selects the best available device:

  - Apple Silicon (M1/M2/M3/M4): uses MPS

  - NVIDIA GPU: uses CUDA

  - Fallback: CPU



Usage:

    from src.segmentation.trainer import SegmentationTrainer



    trainer = SegmentationTrainer(dataset_yaml="data/yolo_dataset/dataset.yaml")

    trainer.train()

    trainer.export()

"""

import os
import platform
from pathlib import Path
from typing import Optional

import torch


# ── Device detection ──────────────────────────────────────────────────────────

def get_best_device() -> str:
    """

    Returns the best available device string for PyTorch / Ultralytics.

    Priority: MPS (Apple Silicon) > CUDA (if compatible) > CPU

    """
    if torch.backends.mps.is_available():
        print("Device: Apple Silicon MPS (GPU accelerated)")
        return "mps"
    elif torch.cuda.is_available():
        # Check CUDA compute capability — RTX 5060 (sm_120) not supported by older PyTorch
        major, minor = torch.cuda.get_device_capability(0)
        sm = major * 10 + minor
        supported = [50, 60, 61, 70, 75, 80, 86, 90]
        if sm in supported:
            gpu = torch.cuda.get_device_name(0)
            print(f"Device: CUDA — {gpu}")
            return "0"
        else:
            gpu = torch.cuda.get_device_name(0)
            print(f"Device: CPU (GPU {gpu} sm_{sm} not supported by this PyTorch build)")
            return "cpu"
    else:
        print("Device: CPU (training will be slow — consider Colab)")
        return "cpu"


# ── Trainer ───────────────────────────────────────────────────────────────────

class SegmentationTrainer:
    """

    Wraps Ultralytics YOLOv8 for floor plan segmentation training.



    Args:

        dataset_yaml:  Path to the dataset.yaml generated by CubiCasaDataset.

        model_size:    YOLOv8 model variant: 'n' (nano), 's' (small), 'm' (medium).

                       For M4 MacBook Air, 's' is the sweet spot.

        output_dir:    Where to save training runs.

        epochs:        Number of training epochs.

        batch_size:    Batch size. Use 8-16 for MPS, 4-8 for CPU.

        img_size:      Input image size for training.

        device:        Override device ('mps', 'cpu', '0'). Auto-detected if None.

    """

    # Model weight files (downloaded automatically on first run)
    MODEL_WEIGHTS = {
        "n": "yolov8n-seg.pt",   # 3.4M params — fastest
        "s": "yolov8s-seg.pt",   # 11.8M params — recommended for M4
        "m": "yolov8m-seg.pt",   # 27.3M params — most accurate
    }

    def __init__(

        self,

        dataset_yaml: str,

        model_size: str = "s",

        output_dir: str = "models/segmentation",

        epochs: int = 100,

        batch_size: int = 8,

        img_size: int = 640,

        device: Optional[str] = None,

        pretrained: bool = True,

    ):
        self.dataset_yaml = Path(dataset_yaml)
        self.model_size = model_size
        self.output_dir = Path(output_dir)
        self.epochs = epochs
        self.batch_size = batch_size
        self.img_size = img_size
        self.device = device or get_best_device()
        self.pretrained = pretrained
        self.model = None

        if model_size not in self.MODEL_WEIGHTS:
            raise ValueError(
                f"model_size must be one of {list(self.MODEL_WEIGHTS.keys())}"
            )
        if not self.dataset_yaml.exists():
            raise FileNotFoundError(f"Dataset YAML not found: {dataset_yaml}")

    def train(self) -> str:
        """

        Start YOLOv8 training.



        Returns:

            Path to the best model weights file.

        """
        from ultralytics import YOLO

        weights = self.MODEL_WEIGHTS[self.model_size]
        print(f"\nStarting training:")
        print(f"  Model:   YOLOv8{self.model_size}-seg ({weights})")
        print(f"  Dataset: {self.dataset_yaml}")
        print(f"  Device:  {self.device}")
        print(f"  Epochs:  {self.epochs}")
        print(f"  Batch:   {self.batch_size}")
        print(f"  Img size:{self.img_size}")
        print()

        self.model = YOLO(weights)

        results = self.model.train(
            data=str(self.dataset_yaml),
            epochs=self.epochs,
            batch=self.batch_size,
            imgsz=self.img_size,
            device=self.device,
            project=str(self.output_dir),
            name="floor_plan_seg",
            # Augmentation — essential for floor plans
            augment=True,
            degrees=10.0,       # Random rotation ±10°
            translate=0.1,      # Random translation
            scale=0.5,          # Random scale
            fliplr=0.5,         # Horizontal flip
            flipud=0.0,         # No vertical flip (plans are top-down)
            # Optimizer settings
            optimizer="AdamW",
            lr0=0.001,
            lrf=0.01,
            weight_decay=0.0005,
            warmup_epochs=3,
            # Save settings
            save=True,
            save_period=10,     # Save checkpoint every 10 epochs
            patience=30,        # Early stopping patience
            # Logging
            verbose=True,
            plots=True,
        )

        best_weights = (
            Path(results.save_dir) / "weights" / "best.pt"
        )
        print(f"\nTraining complete. Best weights: {best_weights}")
        return str(best_weights)

    def resume(self, checkpoint_path: str) -> str:
        """Resume training from a checkpoint."""
        from ultralytics import YOLO
        print(f"Resuming from checkpoint: {checkpoint_path}")
        self.model = YOLO(checkpoint_path)
        results = self.model.train(resume=True)
        return str(Path(results.save_dir) / "weights" / "best.pt")

    def validate(self, weights_path: Optional[str] = None) -> dict:
        """

        Run validation on the test set.



        Returns:

            Dictionary of evaluation metrics (mAP50, mAP50-95, etc.)

        """
        from ultralytics import YOLO

        if weights_path:
            model = YOLO(weights_path)
        elif self.model:
            model = self.model
        else:
            raise ValueError(
                "No model loaded. Run train() first or pass weights_path."
            )

        metrics = model.val(
            data=str(self.dataset_yaml),
            device=self.device,
            split="test",
        )

        print("\nValidation results:")
        print(f"  mAP50:    {metrics.seg.map50:.4f}")
        print(f"  mAP50-95: {metrics.seg.map:.4f}")
        return {
            "map50":    metrics.seg.map50,
            "map50_95": metrics.seg.map,
        }

    def export(

        self,

        weights_path: Optional[str] = None,

        format: str = "onnx",

    ) -> str:
        """

        Export trained model for production deployment.



        Args:

            weights_path: Path to best.pt. Uses last trained model if None.

            format:       Export format: 'onnx', 'coreml' (Apple), 'torchscript'.



        Returns:

            Path to exported model file.

        """
        from ultralytics import YOLO

        if weights_path:
            model = YOLO(weights_path)
        elif self.model:
            model = self.model
        else:
            raise ValueError("No model to export. Run train() first.")

        print(f"Exporting model to {format.upper()} format...")

        # CoreML is best for on-device M4 inference
        if format == "coreml":
            path = model.export(
                format="coreml",
                imgsz=self.img_size,
                nms=True,
            )
        else:
            path = model.export(
                format=format,
                imgsz=self.img_size,
                dynamic=True,
                simplify=True,
            )

        print(f"Exported: {path}")
        return str(path)