train.py

#!/usr/bin/env python3
"""Training script for semantic segmentation with GUI."""

import argparse
import traceback
import yaml
from pathlib import Path
from typing import Dict

import torch
from torch.utils.data import DataLoader
from torch.nn import CrossEntropyLoss
from torch.optim import SGD
from torch.optim.lr_scheduler import CosineAnnealingLR
from super_gradients.training import models

from training import SegmentationDataset
from training.trainer import SegmentationTrainer
from training.training_gui import TrainingGUI
from file_ops import FileManager


def load_config(config_path: Path) -> Dict:
    """Load configuration from YAML file."""
    with open(config_path, 'r') as f:
        config = yaml.safe_load(f)
    return config


def main():
    """Main training function."""
    print("=" * 60)
    print("Starting training with GUI...")
    print("=" * 60)
    
    parser = argparse.ArgumentParser(
        description="Train semantic segmentation model with GUI"
    )
    parser.add_argument(
        "--data_dir",
        type=str,
        default="data",
        help="Base data directory"
    )
    parser.add_argument(
        "--config",
        type=str,
        default="configs/ontology.yaml",
        help="Path to ontology configuration file"
    )
    parser.add_argument(
        "--model",
        type=str,
        default="ddrnet_39",
        help="Model architecture"
    )
    parser.add_argument(
        "--batch_size",
        type=int,
        default=4,
        help="Batch size"
    )
    parser.add_argument(
        "--epochs",
        type=int,
        default=50,
        help="Number of training epochs"
    )
    parser.add_argument(
        "--input_size",
        type=int,
        default=512,
        help="Input image size (square)"
    )
    parser.add_argument(
        "--lr",
        type=float,
        default=0.01,
        help="Learning rate"
    )
    parser.add_argument(
        "--train_split",
        type=float,
        default=0.90,
        help="Training split ratio"
    )
    parser.add_argument(
        "--val_split",
        type=float,
        default=0.10,
        help="Validation split ratio"
    )
    parser.add_argument(
        "--test_split",
        type=float,
        default=0.00,
        help="Test split ratio"
    )
    parser.add_argument(
        "--device",
        type=str,
        default="cuda",
        choices=["cuda", "cpu"],
        help="Device to use"
    )
    parser.add_argument(
        "--no-gui",
        action="store_true",
        help="Disable GUI"
    )
    parser.add_argument(
        "--no-augmentation",
        action="store_true",
        help="Disable data augmentation"
    )
    
    args = parser.parse_args()
    
    # Load config
    config_path = Path(args.config)
    if not config_path.exists():
        raise FileNotFoundError(f"Config file not found: {config_path}")
    config = load_config(config_path)
    
    # Get class info
    class_names = [cls['name'] for cls in config['ontology']['classes']]
    class_colors = [cls['color'] for cls in config['ontology']['classes']]
    num_classes = len(class_names) + 1  # +1 for background
    
    # Setup paths
    script_dir = Path(__file__).parent
    data_dir = script_dir / args.data_dir
    training_dir = data_dir / "training"
    checkpoint_dir = script_dir / "models" / "checkpoints" / f"{args.model}_custom"
    
    print(f"\nLoading data from: {data_dir / 'labeling' / 'accepted'}")
    print(f"Number of classes: {num_classes} (including background)")
    
    # Initialize file manager and prepare data splits
    file_manager = FileManager(data_dir)
    train_images_dir, train_masks_dir, val_images_dir, val_masks_dir = file_manager.prepare_training_data(
        training_dir,
        args.train_split,
        args.val_split,
        args.test_split
    )
    
    # Create datasets
    print("\nCreating datasets...")
    use_aug = not args.no_augmentation
    if use_aug:
        print("✓ Data augmentation enabled")
    else:
        print("✗ Data augmentation disabled")
    
    train_dataset = SegmentationDataset(
        images_dir=train_images_dir,
        masks_dir=train_masks_dir,
        num_classes=num_classes,
        input_size=(args.input_size, args.input_size),
        is_training=True,
        use_augmentation=use_aug
    )
    
    val_dataset = SegmentationDataset(
        images_dir=val_images_dir,
        masks_dir=val_masks_dir,
        num_classes=num_classes,
        input_size=(args.input_size, args.input_size),
        is_training=False,
        use_augmentation=False
    )
    
    print(f"Found {len(train_dataset)} training samples")
    print(f"Found {len(val_dataset)} validation samples")
    
    # Create dataloaders
    train_loader = DataLoader(
        train_dataset,
        batch_size=args.batch_size,
        shuffle=True,
        num_workers=2,
        pin_memory=True
    )
    
    val_loader = DataLoader(
        val_dataset,
        batch_size=args.batch_size,
        shuffle=False,
        num_workers=2,
        pin_memory=True
    )
    
    # Setup device
    device = args.device
    if device == "cuda" and not torch.cuda.is_available():
        print("Warning: CUDA not available, using CPU")
        device = "cpu"
    print(f"\nUsing device: {device}")
    
    # Load model
    print(f"\nInitializing {args.model} model...")
    print("This may take a moment to download pretrained weights...")
    try:
        model = models.get(args.model, num_classes=num_classes)
        print("✓ Model loaded successfully")
    except Exception as e:
        print(f"Error loading model: {e}")
        traceback.print_exc()
        return
    
    # Initialize trainer
    trainer = SegmentationTrainer(
        model=model,
        train_loader=train_loader,
        val_loader=val_loader,
        device=device,
        checkpoint_dir=checkpoint_dir
    )
    
    # Setup optimizer and loss
    optimizer = SGD(
        model.parameters(),
        lr=args.lr,
        momentum=0.9,
        weight_decay=1e-4
    )
    
    criterion = CrossEntropyLoss()
    
    scheduler = CosineAnnealingLR(optimizer, T_max=args.epochs)
    
    # Setup GUI if enabled
    gui = None
    if not args.no_gui:
        print("\nInitializing training GUI...")
        gui = TrainingGUI(
            class_names=class_names,
            class_colors=class_colors,
            num_epochs=args.epochs
        )
        gui.start()
    
    # Training loop
    print(f"\nStarting training for {args.epochs} epochs...")
    print("=" * 60)
    
    try:
        for epoch in range(args.epochs):
            # Train one epoch
            train_loss, train_miou = trainer.train_epoch(optimizer, criterion, num_classes)
            
            # Validate
            val_loss, val_miou = trainer.validate(criterion, num_classes)
            
            # Update scheduler
            if scheduler:
                scheduler.step()
            
            # Save best model
            if val_loss < trainer.best_val_loss:
                trainer.best_val_loss = val_loss
                trainer.save_checkpoint("best_model.pth")
            
            # Save latest
            trainer.save_checkpoint("latest_model.pth")
            
            # Get random samples for visualization
            train_samples, val_samples = trainer.get_random_samples(num_samples=3)
            
            # Update GUI
            if gui and gui.is_running:
                gui.on_epoch_end(
                    epoch=epoch + 1,
                    train_loss=train_loss,
                    val_loss=val_loss,
                    train_miou=train_miou,
                    val_miou=val_miou,
                    train_samples=train_samples,
                    val_samples=val_samples
                )
                gui.update()
            else:
                # Console output
                print(f"Epoch {epoch+1}/{args.epochs}: "
                      f"Train Loss={train_loss:.4f}, Val Loss={val_loss:.4f}, "
                      f"Train mIoU={train_miou:.4f}, Val mIoU={val_miou:.4f}")
    
    except KeyboardInterrupt:
        print("\n\nTraining interrupted by user")
    except Exception as e:
        print(f"\n\nError during training: {e}")
        traceback.print_exc()
    finally:
        if gui:
            print("\nClosing GUI...")
            gui.log("Training finished!")
            gui.update_status("Training complete")
            input("Press Enter to close GUI...")
            gui.close()
    
    print("\n" + "=" * 60)
    print("Training completed!")
    print(f"Checkpoints saved to: {checkpoint_dir}")
    print("=" * 60)


if __name__ == "__main__":
    main()