geoai-py 0.5.6__py2.py3-none-any.whl → 0.7.0__py2.py3-none-any.whl

geoai/train.py

@@ -24,6 +24,16 @@ from tqdm import tqdm
 from .utils import download_model_from_hf
 
 
+# Additional imports for semantic segmentation
+try:
+    import segmentation_models_pytorch as smp
+    from torch.nn import functional as F
+
+    SMP_AVAILABLE = True
+except ImportError:
+    SMP_AVAILABLE = False
+
+
 def get_instance_segmentation_model(num_classes=2, num_channels=3, pretrained=True):
     """
     Get Mask R-CNN model with custom input channels and output classes.
@@ -577,6 +587,7 @@ def train_MaskRCNN_model(
     labels_dir,
     output_dir,
     num_channels=3,
+    model=None,
     pretrained=True,
     pretrained_model_path=None,
     batch_size=4,
@@ -601,6 +612,7 @@ def train_MaskRCNN_model(
         output_dir (str): Directory to save model checkpoints and results.
         num_channels (int, optional): Number of input channels. If None, auto-detected.
             Defaults to 3.
+        model (torch.nn.Module, optional): Predefined model. If None, a new model is created.
         pretrained (bool): Whether to use pretrained backbone. This is ignored if
             pretrained_model_path is provided. Defaults to True.
         pretrained_model_path (str, optional): Path to a .pth file to load as a
@@ -708,9 +720,10 @@ def train_MaskRCNN_model(
     )
 
     # Initialize model (2 classes: background and building)
-    model = get_instance_segmentation_model(
-        num_classes=2, num_channels=num_channels, pretrained=pretrained
-    )
+    if model is None:
+        model = get_instance_segmentation_model(
+            num_classes=2, num_channels=num_channels, pretrained=pretrained
+        )
     model.to(device)
 
     # Set up optimizer
@@ -1088,6 +1101,7 @@ def object_detection(
     confidence_threshold=0.5,
     batch_size=4,
     num_channels=3,
+    model=None,
     pretrained=True,
     device=None,
     **kwargs,
@@ -1104,6 +1118,7 @@ def object_detection(
         confidence_threshold (float): Confidence threshold for predictions (0-1).
         batch_size (int): Batch size for inference.
         num_channels (int): Number of channels in the input image and model.
+        model (torch.nn.Module, optional): Predefined model. If None, a new model is created.
         pretrained (bool): Whether to use pretrained backbone for model loading.
         device (torch.device, optional): Device to run inference on. If None, uses CUDA if available.
         **kwargs: Additional arguments passed to inference_on_geotiff.
@@ -1116,9 +1131,10 @@ def object_detection(
         device = (
             torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
         )
-    model = get_instance_segmentation_model(
-        num_classes=2, num_channels=num_channels, pretrained=pretrained
-    )
+    if model is None:
+        model = get_instance_segmentation_model(
+            num_classes=2, num_channels=num_channels, pretrained=pretrained
+        )
 
     if not os.path.exists(model_path):
         try:
@@ -1153,6 +1169,7 @@ def object_detection_batch(
     overlap=256,
     confidence_threshold=0.5,
     batch_size=4,
+    model=None,
     num_channels=3,
     pretrained=True,
     device=None,
@@ -1174,6 +1191,7 @@ def object_detection_batch(
         confidence_threshold (float): Confidence threshold for predictions (0-1).
         batch_size (int): Batch size for inference.
         num_channels (int): Number of channels in the input image and model.
+        model (torch.nn.Module, optional): Predefined model. If None, a new model is created.
         pretrained (bool): Whether to use pretrained backbone for model loading.
         device (torch.device, optional): Device to run inference on. If None, uses CUDA if available.
         **kwargs: Additional arguments passed to inference_on_geotiff.
@@ -1186,9 +1204,10 @@ def object_detection_batch(
         device = (
             torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
         )
-    model = get_instance_segmentation_model(
-        num_classes=2, num_channels=num_channels, pretrained=pretrained
-    )
+    if model is None:
+        model = get_instance_segmentation_model(
+            num_classes=2, num_channels=num_channels, pretrained=pretrained
+        )
 
     if not os.path.exists(output_dir):
         os.makedirs(output_dir, exist_ok=True)
@@ -1232,3 +1251,1016 @@ def object_detection_batch(
             device=device,
             **kwargs,
         )
+
+
+class SemanticSegmentationDataset(Dataset):
+    """Dataset for semantic segmentation from GeoTIFF images and labels."""
+
+    def __init__(self, image_paths, label_paths, transforms=None, num_channels=None):
+        """
+        Initialize dataset for semantic segmentation.
+
+        Args:
+            image_paths (list): List of paths to image GeoTIFF files.
+            label_paths (list): List of paths to label GeoTIFF files.
+            transforms (callable, optional): Transformations to apply to images and masks.
+            num_channels (int, optional): Number of channels to use from images. If None,
+                auto-detected from the first image.
+        """
+        self.image_paths = image_paths
+        self.label_paths = label_paths
+        self.transforms = transforms
+
+        # Auto-detect the number of channels if not specified
+        if num_channels is None:
+            with rasterio.open(self.image_paths[0]) as src:
+                self.num_channels = src.count
+        else:
+            self.num_channels = num_channels
+
+    def __len__(self):
+        return len(self.image_paths)
+
+    def __getitem__(self, idx):
+        # Load image
+        with rasterio.open(self.image_paths[idx]) as src:
+            # Read as [C, H, W] format
+            image = src.read().astype(np.float32)
+
+            # Normalize image to [0, 1] range
+            image = image / 255.0
+
+            # Handle different number of channels
+            if image.shape[0] > self.num_channels:
+                image = image[: self.num_channels]  # Keep only specified bands
+            elif image.shape[0] < self.num_channels:
+                # Pad with zeros if less than specified bands
+                padded = np.zeros(
+                    (self.num_channels, image.shape[1], image.shape[2]),
+                    dtype=np.float32,
+                )
+                padded[: image.shape[0]] = image
+                image = padded
+
+            # Convert to CHW tensor
+            image = torch.as_tensor(image, dtype=torch.float32)
+
+        # Load label mask
+        with rasterio.open(self.label_paths[idx]) as src:
+            label_mask = src.read(1).astype(np.int64)
+            # Keep original class values for multi-class segmentation
+            # No conversion to binary - preserve all class labels
+
+        # Convert to tensor
+        mask = torch.as_tensor(label_mask, dtype=torch.long)
+
+        # Apply transforms if specified
+        if self.transforms is not None:
+            image, mask = self.transforms(image, mask)
+
+        return image, mask
+
+
+class SemanticTransforms:
+    """Custom transforms for semantic segmentation."""
+
+    def __init__(self, transforms):
+        self.transforms = transforms
+
+    def __call__(self, image, mask):
+        for t in self.transforms:
+            image, mask = t(image, mask)
+        return image, mask
+
+
+class SemanticToTensor:
+    """Convert numpy.ndarray to tensor for semantic segmentation."""
+
+    def __call__(self, image, mask):
+        return image, mask
+
+
+class SemanticRandomHorizontalFlip:
+    """Random horizontal flip transform for semantic segmentation."""
+
+    def __init__(self, prob=0.5):
+        self.prob = prob
+
+    def __call__(self, image, mask):
+        if random.random() < self.prob:
+            # Flip image and mask along width dimension
+            image = torch.flip(image, dims=[2])
+            mask = torch.flip(mask, dims=[1])
+        return image, mask
+
+
+def get_semantic_transform(train):
+    """
+    Get transforms for semantic segmentation data augmentation.
+
+    Args:
+        train (bool): Whether to include training-specific transforms.
+
+    Returns:
+        SemanticTransforms: Composed transforms.
+    """
+    transforms = []
+    transforms.append(SemanticToTensor())
+
+    if train:
+        transforms.append(SemanticRandomHorizontalFlip(0.5))
+
+    return SemanticTransforms(transforms)
+
+
+def get_smp_model(
+    architecture="unet",
+    encoder_name="resnet34",
+    encoder_weights="imagenet",
+    in_channels=3,
+    classes=2,
+    activation=None,
+    **kwargs,
+):
+    """
+    Get a segmentation model from segmentation-models-pytorch using the generic create_model function.
+
+    Args:
+        architecture (str): Model architecture (e.g., 'unet', 'deeplabv3', 'deeplabv3plus', 'fpn',
+            'pspnet', 'linknet', 'manet', 'pan', 'upernet', etc.). Case insensitive.
+        encoder_name (str): Encoder backbone name (e.g., 'resnet34', 'efficientnet-b0', 'mit_b0', etc.).
+        encoder_weights (str): Encoder weights ('imagenet' or None).
+        in_channels (int): Number of input channels.
+        classes (int): Number of output classes.
+        activation (str): Activation function for output layer.
+        **kwargs: Additional arguments passed to smp.create_model().
+
+    Returns:
+        torch.nn.Module: Segmentation model.
+
+    Note:
+        This function uses smp.create_model() which supports all architectures available in
+        segmentation-models-pytorch, making it future-proof for new model additions.
+    """
+    if not SMP_AVAILABLE:
+        raise ImportError(
+            "segmentation-models-pytorch is not installed. "
+            "Please install it with: pip install segmentation-models-pytorch"
+        )
+
+    try:
+        # Use the generic create_model function - supports all SMP architectures
+        model = smp.create_model(
+            arch=architecture,  # Case insensitive
+            encoder_name=encoder_name,
+            encoder_weights=encoder_weights,
+            in_channels=in_channels,
+            classes=classes,
+            **kwargs,
+        )
+
+        # Apply activation if specified (note: activation is handled differently in create_model)
+        if activation is not None:
+            import warnings
+
+            warnings.warn(
+                "The 'activation' parameter is deprecated when using smp.create_model(). "
+                "Apply activation manually after model creation if needed.",
+                DeprecationWarning,
+                stacklevel=2,
+            )
+
+        return model
+
+    except Exception as e:
+        # Provide helpful error message
+        available_archs = []
+        try:
+            # Try to get available architectures from smp
+            if hasattr(smp, "get_available_models"):
+                available_archs = smp.get_available_models()
+            else:
+                available_archs = [
+                    "unet",
+                    "unetplusplus",
+                    "manet",
+                    "linknet",
+                    "fpn",
+                    "pspnet",
+                    "deeplabv3",
+                    "deeplabv3plus",
+                    "pan",
+                    "upernet",
+                ]
+        except:
+            available_archs = [
+                "unet",
+                "fpn",
+                "deeplabv3plus",
+                "pspnet",
+                "linknet",
+                "manet",
+            ]
+
+        raise ValueError(
+            f"Failed to create model with architecture '{architecture}' and encoder '{encoder_name}'. "
+            f"Error: {str(e)}. "
+            f"Available architectures include: {', '.join(available_archs)}. "
+            f"Please check the segmentation-models-pytorch documentation for supported combinations."
+        )
+
+
+def dice_coefficient(pred, target, smooth=1e-6, num_classes=None):
+    """
+    Calculate Dice coefficient for segmentation (binary or multi-class).
+
+    Args:
+        pred (torch.Tensor): Predicted mask (probabilities or logits) with shape [C, H, W] or [H, W].
+        target (torch.Tensor): Ground truth mask with shape [H, W].
+        smooth (float): Smoothing factor to avoid division by zero.
+        num_classes (int, optional): Number of classes. If None, auto-detected.
+
+    Returns:
+        float: Mean Dice coefficient across all classes.
+    """
+    # Convert predictions to class predictions
+    if pred.dim() == 3:  # [C, H, W] format
+        pred = torch.softmax(pred, dim=0)
+        pred_classes = torch.argmax(pred, dim=0)
+    elif pred.dim() == 2:  # [H, W] format
+        pred_classes = pred
+    else:
+        raise ValueError(f"Unexpected prediction dimensions: {pred.shape}")
+
+    # Auto-detect number of classes if not provided
+    if num_classes is None:
+        num_classes = max(pred_classes.max().item(), target.max().item()) + 1
+
+    # Calculate Dice for each class and average
+    dice_scores = []
+    for class_id in range(num_classes):
+        pred_class = (pred_classes == class_id).float()
+        target_class = (target == class_id).float()
+
+        intersection = (pred_class * target_class).sum()
+        union = pred_class.sum() + target_class.sum()
+
+        if union > 0:
+            dice = (2.0 * intersection + smooth) / (union + smooth)
+            dice_scores.append(dice.item())
+
+    return sum(dice_scores) / len(dice_scores) if dice_scores else 0.0
+
+
+def iou_coefficient(pred, target, smooth=1e-6, num_classes=None):
+    """
+    Calculate IoU coefficient for segmentation (binary or multi-class).
+
+    Args:
+        pred (torch.Tensor): Predicted mask (probabilities or logits) with shape [C, H, W] or [H, W].
+        target (torch.Tensor): Ground truth mask with shape [H, W].
+        smooth (float): Smoothing factor to avoid division by zero.
+        num_classes (int, optional): Number of classes. If None, auto-detected.
+
+    Returns:
+        float: Mean IoU coefficient across all classes.
+    """
+    # Convert predictions to class predictions
+    if pred.dim() == 3:  # [C, H, W] format
+        pred = torch.softmax(pred, dim=0)
+        pred_classes = torch.argmax(pred, dim=0)
+    elif pred.dim() == 2:  # [H, W] format
+        pred_classes = pred
+    else:
+        raise ValueError(f"Unexpected prediction dimensions: {pred.shape}")
+
+    # Auto-detect number of classes if not provided
+    if num_classes is None:
+        num_classes = max(pred_classes.max().item(), target.max().item()) + 1
+
+    # Calculate IoU for each class and average
+    iou_scores = []
+    for class_id in range(num_classes):
+        pred_class = (pred_classes == class_id).float()
+        target_class = (target == class_id).float()
+
+        intersection = (pred_class * target_class).sum()
+        union = pred_class.sum() + target_class.sum() - intersection
+
+        if union > 0:
+            iou = (intersection + smooth) / (union + smooth)
+            iou_scores.append(iou.item())
+
+    return sum(iou_scores) / len(iou_scores) if iou_scores else 0.0
+
+
+def train_semantic_one_epoch(
+    model, optimizer, data_loader, device, epoch, criterion, print_freq=10, verbose=True
+):
+    """
+    Train the semantic segmentation model for one epoch.
+
+    Args:
+        model (torch.nn.Module): The model to train.
+        optimizer (torch.optim.Optimizer): The optimizer to use.
+        data_loader (torch.utils.data.DataLoader): DataLoader for training data.
+        device (torch.device): Device to train on.
+        epoch (int): Current epoch number.
+        criterion: Loss function.
+        print_freq (int): How often to print progress.
+        verbose (bool): Whether to print detailed progress.
+
+    Returns:
+        float: Average loss for the epoch.
+    """
+    model.train()
+    total_loss = 0
+    num_batches = len(data_loader)
+
+    start_time = time.time()
+
+    for i, (images, targets) in enumerate(data_loader):
+        # Move images and targets to device
+        images = images.to(device)
+        targets = targets.to(device)
+
+        # Forward pass
+        outputs = model(images)
+        loss = criterion(outputs, targets)
+
+        # Backward pass
+        optimizer.zero_grad()
+        loss.backward()
+        optimizer.step()
+
+        # Track loss
+        total_loss += loss.item()
+
+        # Print progress
+        if i % print_freq == 0:
+            elapsed_time = time.time() - start_time
+            if verbose:
+                print(
+                    f"Epoch: {epoch}, Batch: {i}/{num_batches}, Loss: {loss.item():.4f}, Time: {elapsed_time:.2f}s"
+                )
+            start_time = time.time()
+
+    # Calculate average loss
+    avg_loss = total_loss / num_batches
+    return avg_loss
+
+
+def evaluate_semantic(model, data_loader, device, criterion, num_classes=2):
+    """
+    Evaluate the semantic segmentation model on the validation set.
+
+    Args:
+        model (torch.nn.Module): The model to evaluate.
+        data_loader (torch.utils.data.DataLoader): DataLoader for validation data.
+        device (torch.device): Device to evaluate on.
+        criterion: Loss function.
+        num_classes (int): Number of classes for evaluation metrics.
+
+    Returns:
+        dict: Evaluation metrics including loss, IoU, and Dice.
+    """
+    model.eval()
+
+    total_loss = 0
+    dice_scores = []
+    iou_scores = []
+    num_batches = len(data_loader)
+
+    with torch.no_grad():
+        for images, targets in data_loader:
+            # Move to device
+            images = images.to(device)
+            targets = targets.to(device)
+
+            # Forward pass
+            outputs = model(images)
+            loss = criterion(outputs, targets)
+            total_loss += loss.item()
+
+            # Calculate metrics for each sample in the batch
+            for pred, target in zip(outputs, targets):
+                dice = dice_coefficient(pred, target, num_classes=num_classes)
+                iou = iou_coefficient(pred, target, num_classes=num_classes)
+                dice_scores.append(dice)
+                iou_scores.append(iou)
+
+    # Calculate metrics
+    avg_loss = total_loss / num_batches
+    avg_dice = sum(dice_scores) / len(dice_scores) if dice_scores else 0
+    avg_iou = sum(iou_scores) / len(iou_scores) if iou_scores else 0
+
+    return {"loss": avg_loss, "Dice": avg_dice, "IoU": avg_iou}
+
+
+def train_segmentation_model(
+    images_dir,
+    labels_dir,
+    output_dir,
+    architecture="unet",
+    encoder_name="resnet34",
+    encoder_weights="imagenet",
+    num_channels=3,
+    num_classes=2,
+    batch_size=8,
+    num_epochs=50,
+    learning_rate=0.001,
+    weight_decay=1e-4,
+    seed=42,
+    val_split=0.2,
+    print_freq=10,
+    verbose=True,
+    save_best_only=True,
+    plot_curves=False,
+    **kwargs,
+):
+    """
+    Train a semantic segmentation model for object detection using segmentation-models-pytorch.
+
+    This function trains a semantic segmentation model for object detection (e.g., building detection)
+    using models from the segmentation-models-pytorch library. Unlike instance segmentation (Mask R-CNN),
+    this approach treats the task as pixel-level binary classification.
+
+    Args:
+        images_dir (str): Directory containing image GeoTIFF files.
+        labels_dir (str): Directory containing label GeoTIFF files.
+        output_dir (str): Directory to save model checkpoints and results.
+        architecture (str): Model architecture ('unet', 'deeplabv3', 'deeplabv3plus', 'fpn',
+            'pspnet', 'linknet', 'manet'). Defaults to 'unet'.
+        encoder_name (str): Encoder backbone name (e.g., 'resnet34', 'resnet50', 'efficientnet-b0').
+            Defaults to 'resnet34'.
+        encoder_weights (str): Encoder pretrained weights ('imagenet' or None). Defaults to 'imagenet'.
+        num_channels (int): Number of input channels. Defaults to 3.
+        num_classes (int): Number of output classes (typically 2 for binary segmentation). Defaults to 2.
+        batch_size (int): Batch size for training. Defaults to 8.
+        num_epochs (int): Number of training epochs. Defaults to 50.
+        learning_rate (float): Initial learning rate. Defaults to 0.001.
+        weight_decay (float): Weight decay for optimizer. Defaults to 1e-4.
+        seed (int): Random seed for reproducibility. Defaults to 42.
+        val_split (float): Fraction of data to use for validation (0-1). Defaults to 0.2.
+        print_freq (int): Frequency of printing training progress. Defaults to 10.
+        verbose (bool): If True, prints detailed training progress. Defaults to True.
+        save_best_only (bool): If True, only saves the best model. Otherwise saves all checkpoints.
+            Defaults to True.
+        plot_curves (bool): If True, plots training curves. Defaults to False.
+        **kwargs: Additional arguments passed to smp.create_model().
+    Returns:
+        None: Model weights are saved to output_dir.
+
+    Raises:
+        ImportError: If segmentation-models-pytorch is not installed.
+        FileNotFoundError: If input directories don't exist or contain no matching files.
+    """
+    import datetime
+
+    if not SMP_AVAILABLE:
+        raise ImportError(
+            "segmentation-models-pytorch is not installed. "
+            "Please install it with: pip install segmentation-models-pytorch"
+        )
+
+    # Set random seeds for reproducibility
+    torch.manual_seed(seed)
+    np.random.seed(seed)
+    random.seed(seed)
+    torch.backends.cudnn.deterministic = True
+    torch.backends.cudnn.benchmark = False
+
+    # Create output directory
+    os.makedirs(output_dir, exist_ok=True)
+
+    # Get device
+    device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
+    print(f"Using device: {device}")
+
+    # Get all image and label files
+    image_files = sorted(
+        [
+            os.path.join(images_dir, f)
+            for f in os.listdir(images_dir)
+            if f.endswith(".tif")
+        ]
+    )
+    label_files = sorted(
+        [
+            os.path.join(labels_dir, f)
+            for f in os.listdir(labels_dir)
+            if f.endswith(".tif")
+        ]
+    )
+
+    print(f"Found {len(image_files)} image files and {len(label_files)} label files")
+
+    # Ensure matching files
+    if len(image_files) != len(label_files):
+        print("Warning: Number of image files and label files don't match!")
+        # Find matching files by basename
+        basenames = [os.path.basename(f) for f in image_files]
+        label_files = [
+            os.path.join(labels_dir, os.path.basename(f))
+            for f in image_files
+            if os.path.exists(os.path.join(labels_dir, os.path.basename(f)))
+        ]
+        image_files = [
+            f
+            for f, b in zip(image_files, basenames)
+            if os.path.exists(os.path.join(labels_dir, b))
+        ]
+        print(f"Using {len(image_files)} matching files")
+
+    if len(image_files) == 0:
+        raise FileNotFoundError("No matching image and label files found")
+
+    # Split data into train and validation sets
+    train_imgs, val_imgs, train_labels, val_labels = train_test_split(
+        image_files, label_files, test_size=val_split, random_state=seed
+    )
+
+    print(f"Training on {len(train_imgs)} images, validating on {len(val_imgs)} images")
+
+    # Create datasets
+    train_dataset = SemanticSegmentationDataset(
+        train_imgs,
+        train_labels,
+        transforms=get_semantic_transform(train=True),
+        num_channels=num_channels,
+    )
+    val_dataset = SemanticSegmentationDataset(
+        val_imgs,
+        val_labels,
+        transforms=get_semantic_transform(train=False),
+        num_channels=num_channels,
+    )
+
+    # Create data loaders
+    train_loader = DataLoader(
+        train_dataset,
+        batch_size=batch_size,
+        shuffle=True,
+        num_workers=4,
+        pin_memory=True,
+    )
+
+    val_loader = DataLoader(
+        val_dataset,
+        batch_size=batch_size,
+        shuffle=False,
+        num_workers=4,
+        pin_memory=True,
+    )
+
+    # Initialize model
+    model = get_smp_model(
+        architecture=architecture,
+        encoder_name=encoder_name,
+        encoder_weights=encoder_weights,
+        in_channels=num_channels,
+        classes=num_classes,
+        activation=None,  # We'll apply softmax later
+        **kwargs,
+    )
+    model.to(device)
+
+    # Set up loss function (CrossEntropyLoss for multi-class, can also use DiceLoss)
+    criterion = torch.nn.CrossEntropyLoss()
+
+    # Set up optimizer
+    optimizer = torch.optim.Adam(
+        model.parameters(), lr=learning_rate, weight_decay=weight_decay
+    )
+
+    # Set up learning rate scheduler
+    lr_scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
+        optimizer, mode="min", factor=0.5, patience=5, verbose=True
+    )
+
+    # Initialize tracking variables
+    best_iou = 0
+    train_losses = []
+    val_losses = []
+    val_ious = []
+    val_dices = []
+
+    print(f"Starting training with {architecture} + {encoder_name}")
+    print(f"Model parameters: {sum(p.numel() for p in model.parameters()):,}")
+
+    # Training loop
+    for epoch in range(num_epochs):
+        # Train one epoch
+        train_loss = train_semantic_one_epoch(
+            model,
+            optimizer,
+            train_loader,
+            device,
+            epoch,
+            criterion,
+            print_freq,
+            verbose,
+        )
+        train_losses.append(train_loss)
+
+        # Evaluate on validation set
+        eval_metrics = evaluate_semantic(
+            model, val_loader, device, criterion, num_classes=num_classes
+        )
+        val_losses.append(eval_metrics["loss"])
+        val_ious.append(eval_metrics["IoU"])
+        val_dices.append(eval_metrics["Dice"])
+
+        # Update learning rate
+        lr_scheduler.step(eval_metrics["loss"])
+
+        # Print metrics
+        print(
+            f"Epoch {epoch+1}/{num_epochs}: "
+            f"Train Loss: {train_loss:.4f}, "
+            f"Val Loss: {eval_metrics['loss']:.4f}, "
+            f"Val IoU: {eval_metrics['IoU']:.4f}, "
+            f"Val Dice: {eval_metrics['Dice']:.4f}"
+        )
+
+        # Save best model
+        if eval_metrics["IoU"] > best_iou:
+            best_iou = eval_metrics["IoU"]
+            print(f"Saving best model with IoU: {best_iou:.4f}")
+            torch.save(model.state_dict(), os.path.join(output_dir, "best_model.pth"))
+
+        # Save checkpoint every 10 epochs (if not save_best_only)
+        if not save_best_only and ((epoch + 1) % 10 == 0 or epoch == num_epochs - 1):
+            torch.save(
+                {
+                    "epoch": epoch,
+                    "model_state_dict": model.state_dict(),
+                    "optimizer_state_dict": optimizer.state_dict(),
+                    "scheduler_state_dict": lr_scheduler.state_dict(),
+                    "best_iou": best_iou,
+                    "architecture": architecture,
+                    "encoder_name": encoder_name,
+                    "num_channels": num_channels,
+                    "num_classes": num_classes,
+                },
+                os.path.join(output_dir, f"checkpoint_epoch_{epoch+1}.pth"),
+            )
+
+    # Save final model
+    torch.save(model.state_dict(), os.path.join(output_dir, "final_model.pth"))
+
+    # Save training history
+    history = {
+        "train_losses": train_losses,
+        "val_losses": val_losses,
+        "val_ious": val_ious,
+        "val_dices": val_dices,
+    }
+    torch.save(history, os.path.join(output_dir, "training_history.pth"))
+
+    # Save training summary
+    with open(os.path.join(output_dir, "training_summary.txt"), "w") as f:
+        f.write(
+            f"Training completed on: {datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n"
+        )
+        f.write(f"Architecture: {architecture}\n")
+        f.write(f"Encoder: {encoder_name}\n")
+        f.write(f"Total epochs: {num_epochs}\n")
+        f.write(f"Best validation IoU: {best_iou:.4f}\n")
+        f.write(f"Final validation IoU: {val_ious[-1]:.4f}\n")
+        f.write(f"Final validation Dice: {val_dices[-1]:.4f}\n")
+        f.write(f"Final validation loss: {val_losses[-1]:.4f}\n")
+
+    print(f"Training complete! Best IoU: {best_iou:.4f}")
+    print(f"Models saved to {output_dir}")
+
+    # Plot training curves
+    if plot_curves:
+        try:
+            plt.figure(figsize=(15, 5))
+
+            plt.subplot(1, 3, 1)
+            plt.plot(train_losses, label="Train Loss")
+            plt.plot(val_losses, label="Val Loss")
+            plt.title("Loss")
+            plt.xlabel("Epoch")
+            plt.ylabel("Loss")
+            plt.legend()
+            plt.grid(True)
+
+            plt.subplot(1, 3, 2)
+            plt.plot(val_ious, label="Val IoU")
+            plt.title("IoU Score")
+            plt.xlabel("Epoch")
+            plt.ylabel("IoU")
+            plt.legend()
+            plt.grid(True)
+
+            plt.subplot(1, 3, 3)
+            plt.plot(val_dices, label="Val Dice")
+            plt.title("Dice Score")
+            plt.xlabel("Epoch")
+            plt.ylabel("Dice")
+            plt.legend()
+            plt.grid(True)
+
+            plt.tight_layout()
+            plt.savefig(
+                os.path.join(output_dir, "training_curves.png"),
+                dpi=150,
+                bbox_inches="tight",
+            )
+            print(
+                f"Training curves saved to {os.path.join(output_dir, 'training_curves.png')}"
+            )
+            plt.close()
+        except Exception as e:
+            print(f"Could not save training curves: {e}")
+
+
+def semantic_inference_on_geotiff(
+    model,
+    geotiff_path,
+    output_path,
+    window_size=512,
+    overlap=256,
+    batch_size=4,
+    num_channels=3,
+    num_classes=2,
+    device=None,
+    **kwargs,
+):
+    """
+    Perform semantic segmentation inference on a large GeoTIFF using a sliding window approach.
+
+    Args:
+        model (torch.nn.Module): Trained semantic segmentation model.
+        geotiff_path (str): Path to input GeoTIFF file.
+        output_path (str): Path to save output mask GeoTIFF.
+        window_size (int): Size of sliding window for inference.
+        overlap (int): Overlap between adjacent windows.
+        batch_size (int): Batch size for inference.
+        num_channels (int): Number of channels to use from the input image.
+        num_classes (int): Number of classes in the model output.
+        device (torch.device, optional): Device to run inference on.
+        **kwargs: Additional arguments.
+
+    Returns:
+        tuple: Tuple containing output path and inference time in seconds.
+    """
+    if device is None:
+        device = (
+            torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
+        )
+
+    # Put model in evaluation mode
+    model.to(device)
+    model.eval()
+
+    # Open the GeoTIFF
+    with rasterio.open(geotiff_path) as src:
+        # Read metadata
+        meta = src.meta
+        height = src.height
+        width = src.width
+
+        # Update metadata for output raster
+        out_meta = meta.copy()
+        out_meta.update({"count": 1, "dtype": "uint8"})
+
+        # Initialize accumulator arrays for multi-class probability blending
+        # We'll accumulate probabilities for each class and then take argmax
+        prob_accumulator = np.zeros((num_classes, height, width), dtype=np.float32)
+        count_accumulator = np.zeros((height, width), dtype=np.float32)
+
+        # Calculate steps
+        steps_y = math.ceil((height - overlap) / (window_size - overlap))
+        steps_x = math.ceil((width - overlap) / (window_size - overlap))
+        last_y = height - window_size
+        last_x = width - window_size
+
+        total_windows = steps_y * steps_x
+        print(f"Processing {total_windows} windows...")
+
+        pbar = tqdm(total=total_windows)
+        batch_inputs = []
+        batch_positions = []
+        batch_count = 0
+
+        start_time = time.time()
+
+        for i in range(steps_y + 1):
+            y = min(i * (window_size - overlap), last_y)
+            y = max(0, y)
+
+            if y > last_y and i > 0:
+                continue
+
+            for j in range(steps_x + 1):
+                x = min(j * (window_size - overlap), last_x)
+                x = max(0, x)
+
+                if x > last_x and j > 0:
+                    continue
+
+                # Read window
+                window = src.read(window=Window(x, y, window_size, window_size))
+
+                if window.shape[1] == 0 or window.shape[2] == 0:
+                    continue
+
+                current_height = window.shape[1]
+                current_width = window.shape[2]
+
+                # Normalize and prepare input
+                image = window.astype(np.float32) / 255.0
+
+                # Handle different number of bands
+                if image.shape[0] > num_channels:
+                    image = image[:num_channels]
+                elif image.shape[0] < num_channels:
+                    padded = np.zeros(
+                        (num_channels, current_height, current_width), dtype=np.float32
+                    )
+                    padded[: image.shape[0]] = image
+                    image = padded
+
+                # Convert to tensor
+                image_tensor = torch.tensor(image, device=device)
+
+                # Add to batch
+                batch_inputs.append(image_tensor)
+                batch_positions.append((y, x, current_height, current_width))
+                batch_count += 1
+
+                # Process batch
+                if batch_count == batch_size or (i == steps_y and j == steps_x):
+                    with torch.no_grad():
+                        batch_tensor = torch.stack(batch_inputs)
+                        outputs = model(batch_tensor)
+
+                        # Apply softmax to get class probabilities
+                        probs = torch.softmax(outputs, dim=1)
+
+                    # Process each output in the batch
+                    for idx, prob in enumerate(probs):
+                        y_pos, x_pos, h, w = batch_positions[idx]
+
+                        # Create weight matrix for blending
+                        y_grid, x_grid = np.mgrid[0:h, 0:w]
+                        dist_from_left = x_grid
+                        dist_from_right = w - x_grid - 1
+                        dist_from_top = y_grid
+                        dist_from_bottom = h - y_grid - 1
+
+                        edge_distance = np.minimum.reduce(
+                            [
+                                dist_from_left,
+                                dist_from_right,
+                                dist_from_top,
+                                dist_from_bottom,
+                            ]
+                        )
+                        edge_distance = np.minimum(edge_distance, overlap / 2)
+
+                        # Avoid zero weights - use minimum weight of 0.1
+                        weight = np.maximum(edge_distance / (overlap / 2), 0.1)
+
+                        # For non-overlapping windows, use uniform weight
+                        if overlap == 0:
+                            weight = np.ones_like(weight)
+
+                        # Convert probabilities to numpy [C, H, W]
+                        prob_np = prob.cpu().numpy()
+
+                        # Accumulate weighted probabilities for each class
+                        y_slice = slice(y_pos, y_pos + h)
+                        x_slice = slice(x_pos, x_pos + w)
+
+                        # Add weighted probabilities for each class
+                        for class_idx in range(num_classes):
+                            prob_accumulator[class_idx, y_slice, x_slice] += (
+                                prob_np[class_idx] * weight
+                            )
+
+                        # Update weight accumulator
+                        count_accumulator[y_slice, x_slice] += weight
+
+                    # Reset batch
+                    batch_inputs = []
+                    batch_positions = []
+                    batch_count = 0
+                    pbar.update(len(probs))
+
+        pbar.close()
+
+        # Calculate final mask by taking argmax of accumulated probabilities
+        mask = np.zeros((height, width), dtype=np.uint8)
+        valid_pixels = count_accumulator > 0
+
+        if np.any(valid_pixels):
+            # Normalize accumulated probabilities by weights
+            normalized_probs = np.zeros_like(prob_accumulator)
+            for class_idx in range(num_classes):
+                normalized_probs[class_idx, valid_pixels] = (
+                    prob_accumulator[class_idx, valid_pixels]
+                    / count_accumulator[valid_pixels]
+                )
+
+            # Take argmax to get final class predictions
+            mask[valid_pixels] = np.argmax(
+                normalized_probs[:, valid_pixels], axis=0
+            ).astype(np.uint8)
+
+            # Check class distribution in predictions (summary only)
+            unique_classes, class_counts = np.unique(
+                mask[valid_pixels], return_counts=True
+            )
+            bg_ratio = np.sum(mask == 0) / mask.size
+            print(
+                f"Predicted classes: {len(unique_classes)} classes, Background: {bg_ratio:.1%}"
+            )
+
+        inference_time = time.time() - start_time
+        print(f"Inference completed in {inference_time:.2f} seconds")
+
+        # Save output
+        with rasterio.open(output_path, "w", **out_meta) as dst:
+            dst.write(mask, 1)
+
+        print(f"Saved prediction to {output_path}")
+
+        return output_path, inference_time
+
+
+def semantic_segmentation(
+    input_path,
+    output_path,
+    model_path,
+    architecture="unet",
+    encoder_name="resnet34",
+    num_channels=3,
+    num_classes=2,
+    window_size=512,
+    overlap=256,
+    batch_size=4,
+    device=None,
+    **kwargs,
+):
+    """
+    Perform semantic segmentation on a GeoTIFF using a trained model.
+
+    Args:
+        input_path (str): Path to input GeoTIFF file.
+        output_path (str): Path to save output mask GeoTIFF.
+        model_path (str): Path to trained model weights.
+        architecture (str): Model architecture used for training.
+        encoder_name (str): Encoder backbone name used for training.
+        num_channels (int): Number of channels in the input image and model.
+        num_classes (int): Number of classes in the model.
+        window_size (int): Size of sliding window for inference.
+        overlap (int): Overlap between adjacent windows.
+        batch_size (int): Batch size for inference.
+        device (torch.device, optional): Device to run inference on.
+        **kwargs: Additional arguments.
+
+    Returns:
+        None: Output mask is saved to output_path.
+    """
+    if device is None:
+        device = (
+            torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
+        )
+
+    # Load model
+    model = get_smp_model(
+        architecture=architecture,
+        encoder_name=encoder_name,
+        encoder_weights=None,  # We're loading trained weights
+        in_channels=num_channels,
+        classes=num_classes,
+        activation=None,
+    )
+
+    if not os.path.exists(model_path):
+        try:
+            model_path = download_model_from_hf(model_path)
+        except Exception as e:
+            raise FileNotFoundError(f"Model file not found: {model_path}")
+
+    model.load_state_dict(torch.load(model_path, map_location=device))
+    model.to(device)
+    model.eval()
+
+    semantic_inference_on_geotiff(
+        model=model,
+        geotiff_path=input_path,
+        output_path=output_path,
+        window_size=window_size,
+        overlap=overlap,
+        batch_size=batch_size,
+        num_channels=num_channels,
+        num_classes=num_classes,
+        device=device,
+        **kwargs,
+    )