""" 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)