landcover_train module

Landcover Classification Training Module

This module extends the base geoai training functionality with specialized components for discrete landcover classification, including: - Enhanced loss functions with boundary weighting - Per-class frequency weighting for imbalanced datasets - Configurable ignore_index handling - Additional validation metrics

Key Features: - Maintains full compatibility with base geoai workflow - Adds optional advanced loss computation modes - Provides flexible ignore_index configuration - Optimized for multi-class landcover segmentation

Author: ValHab Project Date: November 2025

DiceLoss

Bases: Module

Dice loss for semantic segmentation.

Computes the Sørensen–Dice coefficient between predictions and targets, which measures region overlap. Effective for class-imbalanced datasets because it evaluates per-class overlap rather than per-pixel accuracy.

Parameters:

Name	Type	Description	Default
smooth	float	Smoothing constant to avoid division by zero.	1.0
ignore_index	int	Target value that is ignored and does not contribute to the loss.	-100
reduction	str	Reduction to apply: "mean" averages per-class losses, "sum" sums them.	'mean'
weight	Optional[Tensor]	Optional per-class weights tensor of shape (C,).	None

Source code in geoai/landcover_train.py

class DiceLoss(nn.Module):
    """Dice loss for semantic segmentation.

    Computes the Sørensen–Dice coefficient between predictions and targets,
    which measures region overlap. Effective for class-imbalanced datasets
    because it evaluates per-class overlap rather than per-pixel accuracy.

    Args:
        smooth: Smoothing constant to avoid division by zero.
        ignore_index: Target value that is ignored and does not contribute
            to the loss.
        reduction: Reduction to apply: ``"mean"`` averages per-class losses,
            ``"sum"`` sums them.
        weight: Optional per-class weights tensor of shape ``(C,)``.
    """

    def __init__(
        self,
        smooth: float = 1.0,
        ignore_index: int = -100,
        reduction: str = "mean",
        weight: Optional[torch.Tensor] = None,
    ):
        super().__init__()
        self.smooth = smooth
        self.ignore_index = ignore_index
        self.reduction = reduction
        self.weight = weight

    def forward(self, inputs: torch.Tensor, targets: torch.Tensor) -> torch.Tensor:
        """Compute Dice loss.

        Args:
            inputs: Predictions of shape ``(N, C, H, W)`` (logits).
            targets: Ground truth of shape ``(N, H, W)`` with class indices.

        Returns:
            Scalar loss when ``reduction`` is ``"mean"`` or ``"sum"``,
            or a per-class tensor of shape ``(C,)`` when ``"none"``.
        """
        num_classes = inputs.shape[1]
        one_hot, valid_mask = _one_hot_with_ignore(
            targets, num_classes, self.ignore_index
        )
        probs = F.softmax(inputs, dim=1)
        probs = probs * valid_mask.unsqueeze(1)  # zero out ignored pixels

        dims = (0, 2, 3)  # reduce over batch and spatial dims
        intersection = (probs * one_hot).sum(dim=dims)
        cardinality = probs.sum(dim=dims) + one_hot.sum(dim=dims)
        dice = (2.0 * intersection + self.smooth) / (cardinality + self.smooth)
        loss = 1.0 - dice  # (C,)

        if self.weight is not None:
            w = self.weight.to(loss.device)
            loss = loss * w

        if self.reduction == "mean":
            return loss.mean()
        elif self.reduction == "sum":
            return loss.sum()
        elif self.reduction == "none":
            return loss
        raise ValueError(f"Unknown reduction: {self.reduction!r}")

forward(inputs, targets)

Compute Dice loss.

Parameters:

Name	Type	Description	Default
inputs	Tensor	Predictions of shape (N, C, H, W) (logits).	required
targets	Tensor	Ground truth of shape (N, H, W) with class indices.	required

Returns:

Type	Description
Tensor	Scalar loss when reduction is "mean" or "sum",
Tensor	or a per-class tensor of shape (C,) when "none".

Source code in geoai/landcover_train.py

def forward(self, inputs: torch.Tensor, targets: torch.Tensor) -> torch.Tensor:
    """Compute Dice loss.

    Args:
        inputs: Predictions of shape ``(N, C, H, W)`` (logits).
        targets: Ground truth of shape ``(N, H, W)`` with class indices.

    Returns:
        Scalar loss when ``reduction`` is ``"mean"`` or ``"sum"``,
        or a per-class tensor of shape ``(C,)`` when ``"none"``.
    """
    num_classes = inputs.shape[1]
    one_hot, valid_mask = _one_hot_with_ignore(
        targets, num_classes, self.ignore_index
    )
    probs = F.softmax(inputs, dim=1)
    probs = probs * valid_mask.unsqueeze(1)  # zero out ignored pixels

    dims = (0, 2, 3)  # reduce over batch and spatial dims
    intersection = (probs * one_hot).sum(dim=dims)
    cardinality = probs.sum(dim=dims) + one_hot.sum(dim=dims)
    dice = (2.0 * intersection + self.smooth) / (cardinality + self.smooth)
    loss = 1.0 - dice  # (C,)

    if self.weight is not None:
        w = self.weight.to(loss.device)
        loss = loss * w

    if self.reduction == "mean":
        return loss.mean()
    elif self.reduction == "sum":
        return loss.sum()
    elif self.reduction == "none":
        return loss
    raise ValueError(f"Unknown reduction: {self.reduction!r}")

FocalLoss

Bases: Module

Focal Loss for addressing class imbalance in segmentation.

Reference: Lin, T. Y., Goyal, P., Girshick, R., He, K., & Dollár, P. (2017). Focal loss for dense object detection. ICCV.

Parameters:

Name	Type	Description	Default
alpha		Weighting factor in range (0,1) to balance positive/negative examples	1.0
gamma		Exponent of the modulating factor (1 - p_t)^gamma	2.0
ignore_index		Specifies a target value that is ignored	-100
reduction		Specifies the reduction to apply to the output	'mean'
weight		Manual rescaling weight given to each class	None

Source code in geoai/landcover_train.py

class FocalLoss(nn.Module):
    """
    Focal Loss for addressing class imbalance in segmentation.

    Reference: Lin, T. Y., Goyal, P., Girshick, R., He, K., & Dollár, P. (2017).
    Focal loss for dense object detection. ICCV.

    Args:
        alpha: Weighting factor in range (0,1) to balance positive/negative examples
        gamma: Exponent of the modulating factor (1 - p_t)^gamma
        ignore_index: Specifies a target value that is ignored
        reduction: Specifies the reduction to apply to the output
        weight: Manual rescaling weight given to each class
    """

    def __init__(
        self, alpha=1.0, gamma=2.0, ignore_index=-100, reduction="mean", weight=None
    ):
        super(FocalLoss, self).__init__()
        self.alpha = alpha
        self.gamma = gamma
        self.ignore_index = ignore_index
        self.reduction = reduction
        self.weight = weight

    def forward(self, inputs, targets):
        """
        Forward pass of focal loss.

        Args:
            inputs: Predictions (N, C, H, W) where C = number of classes
            targets: Ground truth (N, H, W) with class indices

        Returns:
            Loss value
        """
        # Get class probabilities
        ce_loss = F.cross_entropy(
            inputs,
            targets,
            weight=self.weight,
            ignore_index=self.ignore_index,
            reduction="none",
        )

        # Get probability of true class
        p_t = torch.exp(-ce_loss)

        # Calculate focal loss
        focal_loss = self.alpha * (1 - p_t) ** self.gamma * ce_loss

        # Apply reduction
        if self.reduction == "mean":
            return focal_loss.mean()
        elif self.reduction == "sum":
            return focal_loss.sum()
        else:
            return focal_loss

forward(inputs, targets)

Forward pass of focal loss.

Parameters:

Name	Type	Description	Default
inputs		Predictions (N, C, H, W) where C = number of classes	required
targets		Ground truth (N, H, W) with class indices	required

Returns:

Type	Description
	Loss value

Source code in geoai/landcover_train.py

def forward(self, inputs, targets):
    """
    Forward pass of focal loss.

    Args:
        inputs: Predictions (N, C, H, W) where C = number of classes
        targets: Ground truth (N, H, W) with class indices

    Returns:
        Loss value
    """
    # Get class probabilities
    ce_loss = F.cross_entropy(
        inputs,
        targets,
        weight=self.weight,
        ignore_index=self.ignore_index,
        reduction="none",
    )

    # Get probability of true class
    p_t = torch.exp(-ce_loss)

    # Calculate focal loss
    focal_loss = self.alpha * (1 - p_t) ** self.gamma * ce_loss

    # Apply reduction
    if self.reduction == "mean":
        return focal_loss.mean()
    elif self.reduction == "sum":
        return focal_loss.sum()
    else:
        return focal_loss

LandcoverCrossEntropyLoss

Bases: Module

Enhanced CrossEntropyLoss with optional ignore_index and class weights.

This extends the standard CrossEntropyLoss with more flexible ignore_index handling, specifically designed for landcover classification tasks.

Parameters:

Name	Type	Description	Default
weight	Optional[Tensor]	Manual rescaling weight given to each class	None
ignore_index	Union[int, bool]	Specifies a target value that is ignored. - False: No values ignored (standard behavior) - int: Specific class index to ignore (e.g., 0 for background)	False
reduction	str	Specifies the reduction to apply ('mean', 'sum', 'none')	'mean'

Source code in geoai/landcover_train.py

class LandcoverCrossEntropyLoss(nn.Module):
    """
    Enhanced CrossEntropyLoss with optional ignore_index and class weights.

    This extends the standard CrossEntropyLoss with more flexible ignore_index
    handling, specifically designed for landcover classification tasks.

    Args:
        weight: Manual rescaling weight given to each class
        ignore_index: Specifies a target value that is ignored.
            - False: No values ignored (standard behavior)
            - int: Specific class index to ignore (e.g., 0 for background)
        reduction: Specifies the reduction to apply ('mean', 'sum', 'none')
    """

    def __init__(
        self,
        weight: Optional[torch.Tensor] = None,
        ignore_index: Union[int, bool] = False,
        reduction: str = "mean",
    ):
        super().__init__()
        self.weight = weight
        # Convert ignore_index: int stays as-is, False becomes -100 (PyTorch default)
        self.ignore_index = ignore_index if isinstance(ignore_index, int) else -100
        self.reduction = reduction

    def forward(self, input: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
        """
        Compute cross entropy loss.

        Args:
            input: Predictions (N, C, H, W) where C = number of classes
            target: Ground truth (N, H, W) with class indices

        Returns:
            Loss value
        """
        return F.cross_entropy(
            input,
            target,
            weight=self.weight,
            ignore_index=self.ignore_index,
            reduction=self.reduction,
        )

forward(input, target)

Compute cross entropy loss.

Parameters:

Name	Type	Description	Default
input	Tensor	Predictions (N, C, H, W) where C = number of classes	required
target	Tensor	Ground truth (N, H, W) with class indices	required

Returns:

Type	Description
Tensor	Loss value

Source code in geoai/landcover_train.py

def forward(self, input: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
    """
    Compute cross entropy loss.

    Args:
        input: Predictions (N, C, H, W) where C = number of classes
        target: Ground truth (N, H, W) with class indices

    Returns:
        Loss value
    """
    return F.cross_entropy(
        input,
        target,
        weight=self.weight,
        ignore_index=self.ignore_index,
        reduction=self.reduction,
    )

TverskyLoss

Bases: Module

Tversky loss for semantic segmentation.

Generalises the Dice loss by allowing asymmetric weighting of false positives and false negatives. Setting alpha = beta = 0.5 recovers the standard Dice loss. Increasing beta relative to alpha penalises false negatives more, which improves recall on rare classes.

Reference

Salehi, S. S. M., Erdogmus, D., & Gholipour, A. (2017). Tversky loss function for image segmentation using 3D fully convolutional deep networks. MLMI Workshop, MICCAI.

Parameters:

Name	Type	Description	Default
alpha	float	Weight for false positives.	0.5
beta	float	Weight for false negatives.	0.5
smooth	float	Smoothing constant to avoid division by zero.	1.0
ignore_index	int	Target value that is ignored.	-100
reduction	str	"mean" or "sum" over classes.	'mean'
weight	Optional[Tensor]	Optional per-class weights tensor of shape (C,).	None

Source code in geoai/landcover_train.py

class TverskyLoss(nn.Module):
    """Tversky loss for semantic segmentation.

    Generalises the Dice loss by allowing asymmetric weighting of false
    positives and false negatives.  Setting ``alpha = beta = 0.5`` recovers
    the standard Dice loss.  Increasing ``beta`` relative to ``alpha``
    penalises false negatives more, which improves recall on rare classes.

    Reference:
        Salehi, S. S. M., Erdogmus, D., & Gholipour, A. (2017).
        Tversky loss function for image segmentation using 3D fully
        convolutional deep networks. *MLMI Workshop, MICCAI*.

    Args:
        alpha: Weight for false positives.
        beta: Weight for false negatives.
        smooth: Smoothing constant to avoid division by zero.
        ignore_index: Target value that is ignored.
        reduction: ``"mean"`` or ``"sum"`` over classes.
        weight: Optional per-class weights tensor of shape ``(C,)``.
    """

    def __init__(
        self,
        alpha: float = 0.5,
        beta: float = 0.5,
        smooth: float = 1.0,
        ignore_index: int = -100,
        reduction: str = "mean",
        weight: Optional[torch.Tensor] = None,
    ):
        super().__init__()
        self.alpha = alpha
        self.beta = beta
        self.smooth = smooth
        self.ignore_index = ignore_index
        self.reduction = reduction
        self.weight = weight

    def forward(self, inputs: torch.Tensor, targets: torch.Tensor) -> torch.Tensor:
        """Compute Tversky loss.

        Args:
            inputs: Predictions of shape ``(N, C, H, W)`` (logits).
            targets: Ground truth of shape ``(N, H, W)`` with class indices.

        Returns:
            Scalar loss when ``reduction`` is ``"mean"`` or ``"sum"``,
            or a per-class tensor of shape ``(C,)`` when ``"none"``.
        """
        num_classes = inputs.shape[1]
        one_hot, valid_mask = _one_hot_with_ignore(
            targets, num_classes, self.ignore_index
        )
        probs = F.softmax(inputs, dim=1)
        probs = probs * valid_mask.unsqueeze(1)  # zero out ignored pixels

        dims = (0, 2, 3)
        tp = (probs * one_hot).sum(dim=dims)
        fp = (probs * (1.0 - one_hot)).sum(dim=dims)
        fn = ((1.0 - probs) * one_hot).sum(dim=dims)
        tversky = (tp + self.smooth) / (
            tp + self.alpha * fp + self.beta * fn + self.smooth
        )
        loss = 1.0 - tversky  # (C,)

        if self.weight is not None:
            w = self.weight.to(loss.device)
            loss = loss * w

        if self.reduction == "mean":
            return loss.mean()
        elif self.reduction == "sum":
            return loss.sum()
        elif self.reduction == "none":
            return loss
        raise ValueError(f"Unknown reduction: {self.reduction!r}")

forward(inputs, targets)

Compute Tversky loss.

Parameters:

Name	Type	Description	Default
inputs	Tensor	Predictions of shape (N, C, H, W) (logits).	required
targets	Tensor	Ground truth of shape (N, H, W) with class indices.	required

Returns:

Type	Description
Tensor	Scalar loss when reduction is "mean" or "sum",
Tensor	or a per-class tensor of shape (C,) when "none".

Source code in geoai/landcover_train.py

def forward(self, inputs: torch.Tensor, targets: torch.Tensor) -> torch.Tensor:
    """Compute Tversky loss.

    Args:
        inputs: Predictions of shape ``(N, C, H, W)`` (logits).
        targets: Ground truth of shape ``(N, H, W)`` with class indices.

    Returns:
        Scalar loss when ``reduction`` is ``"mean"`` or ``"sum"``,
        or a per-class tensor of shape ``(C,)`` when ``"none"``.
    """
    num_classes = inputs.shape[1]
    one_hot, valid_mask = _one_hot_with_ignore(
        targets, num_classes, self.ignore_index
    )
    probs = F.softmax(inputs, dim=1)
    probs = probs * valid_mask.unsqueeze(1)  # zero out ignored pixels

    dims = (0, 2, 3)
    tp = (probs * one_hot).sum(dim=dims)
    fp = (probs * (1.0 - one_hot)).sum(dim=dims)
    fn = ((1.0 - probs) * one_hot).sum(dim=dims)
    tversky = (tp + self.smooth) / (
        tp + self.alpha * fp + self.beta * fn + self.smooth
    )
    loss = 1.0 - tversky  # (C,)

    if self.weight is not None:
        w = self.weight.to(loss.device)
        loss = loss * w

    if self.reduction == "mean":
        return loss.mean()
    elif self.reduction == "sum":
        return loss.sum()
    elif self.reduction == "none":
        return loss
    raise ValueError(f"Unknown reduction: {self.reduction!r}")

UnifiedFocalLoss

Bases: Module

Unified Focal Loss combining distribution-based and region-based losses.

Implements the framework from Yeung et al. (2021) which unifies focal cross-entropy (distribution-based) and focal Tversky (region-based) losses into a single compound loss. This is particularly effective for semantic segmentation with severe class imbalance.

The combined loss is::

L = lambda_ * L_dist + (1 - lambda_) * L_region

where L_dist is focal cross-entropy and L_region is focal Tversky.

Reference

Yeung, M., Sala, E., Schönlieb, C.-B., & Rundo, L. (2021). Unified Focal loss: Generalising Dice and cross entropy-based losses to handle class imbalanced medical image segmentation. Computerized Medical Imaging and Graphics, 95, 102026.

Note

Inspired by the implementation in the terrainseg <https://github.com/maxwell-geospatial/terrainseg>_ package by Maxwell (2026).

Parameters:

Name	Type	Description	Default
lambda_	float	Balance between distribution and region components. 1.0 = pure focal CE, 0.0 = pure focal Tversky.	0.5
gamma	float	Focusing parameter for both components. Higher values down-weight easy examples more aggressively.	0.75
delta	float	Tversky false-negative weight; false-positive weight is 1 - delta. Values > 0.5 emphasise recall.	0.6
smooth	float	Smoothing constant for the Tversky denominator.	1.0
ignore_index	int	Target value that is ignored.	-100
weight	Optional[Tensor]	Per-class weights for the distribution (focal CE) component.	None
region_weight	Optional[Tensor]	Per-class weights for the region (focal Tversky) component. Falls back to weight if None.	None
use_log_cosh	bool	Apply log(cosh(loss)) for gradient smoothing.	False

Source code in geoai/landcover_train.py

class UnifiedFocalLoss(nn.Module):
    """Unified Focal Loss combining distribution-based and region-based losses.

    Implements the framework from Yeung et al. (2021) which unifies focal
    cross-entropy (distribution-based) and focal Tversky (region-based) losses
    into a single compound loss.  This is particularly effective for semantic
    segmentation with severe class imbalance.

    The combined loss is::

        L = lambda_ * L_dist + (1 - lambda_) * L_region

    where ``L_dist`` is focal cross-entropy and ``L_region`` is focal Tversky.

    Reference:
        Yeung, M., Sala, E., Schönlieb, C.-B., & Rundo, L. (2021).
        Unified Focal loss: Generalising Dice and cross entropy-based losses
        to handle class imbalanced medical image segmentation.
        *Computerized Medical Imaging and Graphics*, 95, 102026.

    Note:
        Inspired by the implementation in the
        `terrainseg <https://github.com/maxwell-geospatial/terrainseg>`_
        package by Maxwell (2026).

    Args:
        lambda_: Balance between distribution and region components.
            ``1.0`` = pure focal CE, ``0.0`` = pure focal Tversky.
        gamma: Focusing parameter for both components.  Higher values
            down-weight easy examples more aggressively.
        delta: Tversky false-negative weight; false-positive weight is
            ``1 - delta``.  Values > 0.5 emphasise recall.
        smooth: Smoothing constant for the Tversky denominator.
        ignore_index: Target value that is ignored.
        weight: Per-class weights for the distribution (focal CE) component.
        region_weight: Per-class weights for the region (focal Tversky)
            component.  Falls back to ``weight`` if *None*.
        use_log_cosh: Apply ``log(cosh(loss))`` for gradient smoothing.
    """

    def __init__(
        self,
        lambda_: float = 0.5,
        gamma: float = 0.75,
        delta: float = 0.6,
        smooth: float = 1.0,
        ignore_index: int = -100,
        weight: Optional[torch.Tensor] = None,
        region_weight: Optional[torch.Tensor] = None,
        use_log_cosh: bool = False,
    ):
        super().__init__()
        self.lambda_ = lambda_
        self.gamma = gamma
        self.delta = delta
        self.smooth = smooth
        self.ignore_index = ignore_index
        self.weight = weight
        self.region_weight = region_weight
        self.use_log_cosh = use_log_cosh

    def forward(self, inputs: torch.Tensor, targets: torch.Tensor) -> torch.Tensor:
        """Compute unified focal loss.

        Args:
            inputs: Predictions of shape ``(N, C, H, W)`` (logits).
            targets: Ground truth of shape ``(N, H, W)`` with class indices.

        Returns:
            Scalar loss value.
        """
        # --- Distribution component: focal cross-entropy ---
        ce = F.cross_entropy(
            inputs,
            targets,
            weight=self.weight.to(inputs.device) if self.weight is not None else None,
            ignore_index=self.ignore_index,
            reduction="none",
        )
        p_t = torch.exp(-ce)
        focal_ce = ((1.0 - p_t) ** self.gamma) * ce

        # Mean over valid pixels
        valid = targets != self.ignore_index
        if valid.any():
            dist_loss = focal_ce[valid].mean()
        else:
            dist_loss = focal_ce.mean()

        # --- Region component: focal Tversky ---
        num_classes = inputs.shape[1]
        one_hot, valid_mask = _one_hot_with_ignore(
            targets, num_classes, self.ignore_index
        )
        probs = F.softmax(inputs, dim=1)
        probs = probs * valid_mask.unsqueeze(1)  # zero out ignored pixels

        dims = (0, 2, 3)
        tp = (probs * one_hot).sum(dim=dims)
        fp = (probs * (1.0 - one_hot)).sum(dim=dims)
        fn = ((1.0 - probs) * one_hot).sum(dim=dims)
        tversky = (tp + self.smooth) / (
            tp + (1.0 - self.delta) * fp + self.delta * fn + self.smooth
        )
        focal_tversky = (1.0 - tversky) ** self.gamma  # (C,)

        rw = self.region_weight if self.region_weight is not None else self.weight
        if rw is not None:
            rw = rw.to(focal_tversky.device)
            focal_tversky = focal_tversky * rw

        region_loss = focal_tversky.mean()

        # --- Combine ---
        loss = self.lambda_ * dist_loss + (1.0 - self.lambda_) * region_loss

        if self.use_log_cosh:
            # Numerically stable log-cosh: |x| + softplus(-2|x|) - log(2)
            abs_loss = torch.abs(loss)
            loss = abs_loss + F.softplus(-2.0 * abs_loss) - 0.6931471805599453

        return loss

forward(inputs, targets)

Compute unified focal loss.

Parameters:

Name	Type	Description	Default
inputs	Tensor	Predictions of shape (N, C, H, W) (logits).	required
targets	Tensor	Ground truth of shape (N, H, W) with class indices.	required

Returns:

Type	Description
Tensor	Scalar loss value.

Source code in geoai/landcover_train.py

def forward(self, inputs: torch.Tensor, targets: torch.Tensor) -> torch.Tensor:
    """Compute unified focal loss.

    Args:
        inputs: Predictions of shape ``(N, C, H, W)`` (logits).
        targets: Ground truth of shape ``(N, H, W)`` with class indices.

    Returns:
        Scalar loss value.
    """
    # --- Distribution component: focal cross-entropy ---
    ce = F.cross_entropy(
        inputs,
        targets,
        weight=self.weight.to(inputs.device) if self.weight is not None else None,
        ignore_index=self.ignore_index,
        reduction="none",
    )
    p_t = torch.exp(-ce)
    focal_ce = ((1.0 - p_t) ** self.gamma) * ce

    # Mean over valid pixels
    valid = targets != self.ignore_index
    if valid.any():
        dist_loss = focal_ce[valid].mean()
    else:
        dist_loss = focal_ce.mean()

    # --- Region component: focal Tversky ---
    num_classes = inputs.shape[1]
    one_hot, valid_mask = _one_hot_with_ignore(
        targets, num_classes, self.ignore_index
    )
    probs = F.softmax(inputs, dim=1)
    probs = probs * valid_mask.unsqueeze(1)  # zero out ignored pixels

    dims = (0, 2, 3)
    tp = (probs * one_hot).sum(dim=dims)
    fp = (probs * (1.0 - one_hot)).sum(dim=dims)
    fn = ((1.0 - probs) * one_hot).sum(dim=dims)
    tversky = (tp + self.smooth) / (
        tp + (1.0 - self.delta) * fp + self.delta * fn + self.smooth
    )
    focal_tversky = (1.0 - tversky) ** self.gamma  # (C,)

    rw = self.region_weight if self.region_weight is not None else self.weight
    if rw is not None:
        rw = rw.to(focal_tversky.device)
        focal_tversky = focal_tversky * rw

    region_loss = focal_tversky.mean()

    # --- Combine ---
    loss = self.lambda_ * dist_loss + (1.0 - self.lambda_) * region_loss

    if self.use_log_cosh:
        # Numerically stable log-cosh: |x| + softplus(-2|x|) - log(2)
        abs_loss = torch.abs(loss)
        loss = abs_loss + F.softplus(-2.0 * abs_loss) - 0.6931471805599453

    return loss

compute_class_weights(labels_dir, num_classes, ignore_index=-100, custom_multipliers=None, max_weight=50.0, use_inverse_frequency=True)

Compute class weights for imbalanced datasets with optional custom multipliers and maximum weight cap.

Parameters:

Name	Type	Description	Default
labels_dir	str	Directory containing label files	required
num_classes	int	Number of classes	required
ignore_index	Union[int, bool]	Class index to ignore when computing weights. - If int: specific class to ignore (pixels will be excluded from weight calc) - If False: no class ignored (all classes contribute to weights)	-100
custom_multipliers	Optional[Dict[int, float]]	Custom multipliers for specific classes after inverse frequency calculation. Format: {class_id: multiplier} Example: {1: 0.5, 7: 2.0} - reduce class 1 weight by half, double class 7 weight	None
max_weight	float	Maximum allowed weight value to prevent extreme values (default: 50.0)	50.0
use_inverse_frequency	bool	Whether to compute inverse frequency weights. - True (default): Compute inverse frequency weights, then apply custom multipliers - False: Use uniform weights (1.0) for all classes, then apply custom multipliers	True

Returns:

Type	Description
Tensor	Tensor of class weights (num_classes,) with custom adjustments and maximum weight cap applied

Source code in geoai/landcover_train.py

def compute_class_weights(
    labels_dir: str,
    num_classes: int,
    ignore_index: Union[int, bool] = -100,
    custom_multipliers: Optional[Dict[int, float]] = None,
    max_weight: float = 50.0,
    use_inverse_frequency: bool = True,
) -> torch.Tensor:
    """
    Compute class weights for imbalanced datasets with optional custom multipliers and maximum weight cap.

    Args:
        labels_dir: Directory containing label files
        num_classes: Number of classes
        ignore_index: Class index to ignore when computing weights.
            - If int: specific class to ignore (pixels will be excluded from weight calc)
            - If False: no class ignored (all classes contribute to weights)
        custom_multipliers: Custom multipliers for specific classes after inverse frequency calculation.
            Format: {class_id: multiplier}
            Example: {1: 0.5, 7: 2.0} - reduce class 1 weight by half, double class 7 weight
        max_weight: Maximum allowed weight value to prevent extreme values (default: 50.0)
        use_inverse_frequency: Whether to compute inverse frequency weights.
            - True (default): Compute inverse frequency weights, then apply custom multipliers
            - False: Use uniform weights (1.0) for all classes, then apply custom multipliers

    Returns:
        Tensor of class weights (num_classes,) with custom adjustments and maximum weight cap applied
    """
    import os
    import rasterio
    from collections import Counter

    # Count pixels for each class
    class_counts = Counter()
    total_pixels = 0

    # Get all label files
    label_extensions = (".tif", ".tiff", ".png", ".jpg", ".jpeg")
    label_files = [
        os.path.join(labels_dir, f)
        for f in os.listdir(labels_dir)
        if f.lower().endswith(label_extensions)
    ]

    logger.info("Computing class weights from %d label files...", len(label_files))

    for label_file in label_files:
        try:
            with rasterio.open(label_file) as src:
                label_data = src.read(1)
                for class_id in range(num_classes):
                    if isinstance(ignore_index, int) and class_id == ignore_index:
                        continue
                    count = (label_data == class_id).sum()
                    class_counts[class_id] += int(count)
                    total_pixels += int(count)
        except Exception as e:
            logger.warning("Could not read %s: %s", label_file, e)
            continue

    if total_pixels == 0:
        raise ValueError("No valid pixels found in label files")

    # Initialize weights
    weights = torch.ones(num_classes)

    if use_inverse_frequency:
        # Compute inverse frequency weights
        for class_id in range(num_classes):
            if isinstance(ignore_index, int) and class_id == ignore_index:
                weights[class_id] = 0.0
            elif class_counts[class_id] > 0:
                # Inverse frequency: total_pixels / class_pixels
                weights[class_id] = total_pixels / class_counts[class_id]
            else:
                weights[class_id] = 0.0

        # Normalize to have mean weight of 1.0
        non_zero_weights = weights[weights > 0]
        if len(non_zero_weights) > 0:
            weights = weights / non_zero_weights.mean()
    else:
        # Use uniform weights (all 1.0)
        for class_id in range(num_classes):
            if isinstance(ignore_index, int) and class_id == ignore_index:
                weights[class_id] = 0.0

    # Apply custom multipliers if provided
    if custom_multipliers:
        logger.info("Applying custom multipliers: %s", custom_multipliers)
        for class_id, multiplier in custom_multipliers.items():
            if class_id < 0 or class_id >= num_classes:
                logger.warning("Invalid class_id %d, skipping", class_id)
                continue

            original_weight = weights[class_id].item()
            weights[class_id] = weights[class_id] * multiplier
            logger.info(
                "  Class %d: %.4f x %s = %.4f",
                class_id,
                original_weight,
                multiplier,
                weights[class_id].item(),
            )
    else:
        logger.info("No custom multipliers provided, using computed weights as-is")

    # Apply maximum weight cap to prevent extreme values
    weights_capped = False
    logger.info("Applying maximum weight cap of %s...", max_weight)
    for class_id in range(num_classes):
        if weights[class_id] > max_weight:
            logger.info(
                "  Class %d: %.4f -> %s (capped)",
                class_id,
                weights[class_id].item(),
                max_weight,
            )
            weights[class_id] = max_weight
            weights_capped = True

    if not weights_capped:
        logger.info("  No weights exceeded the cap")

    logger.info("Class pixel counts: %s", dict(class_counts))
    logger.info("Final class weights:")
    for class_id in range(num_classes):
        pixel_count = class_counts.get(class_id, 0)
        percent = (pixel_count / total_pixels * 100) if total_pixels > 0 else 0
        logger.info(
            "  Class %d: weight=%.4f, pixels=%s (%.2f%%)",
            class_id,
            weights[class_id].item(),
            f"{pixel_count:,}",
            percent,
        )

    if isinstance(ignore_index, int) and 0 <= ignore_index < num_classes:
        logger.info("Note: Class %d (ignore_index) has weight 0.0", ignore_index)

    return weights

evaluate_sparse_iou(model, images_dir, labels_dir, num_classes, num_channels=3, batch_size=8, background_class=0, ignore_index=False, device=None, verbose=True)

Evaluate a trained model using sparse labels IoU.

This function is designed for incomplete/sparse ground truth where background (0) means "unlabeled" rather than "definitely not this class". Predictions in background areas are NOT penalized as false positives.

Use this for post-training evaluation when your training masks are incomplete.

Parameters:

Name	Type	Description	Default
model	Module	Trained segmentation model	required
images_dir	str	Directory containing validation images	required
labels_dir	str	Directory containing validation labels	required
num_classes	int	Number of classes	required
num_channels	int	Number of input channels (default: 3)	3
batch_size	int	Batch size for evaluation (default: 8)	8
background_class	int	Class ID for background/unlabeled pixels (default: 0)	0
ignore_index	Union[int, bool]	Class to ignore during evaluation. - If int: specific class index to ignore - If False: no class ignored (default)	False
device	Optional[device]	Torch device (auto-detected if None)	None
verbose	bool	Print detailed results (default: True)	True

Returns:

Type	Description
Dict[str, Any]	Dictionary containing:
Dict[str, Any]	'mean_sparse_iou': Mean IoU across all non-background classes
Dict[str, Any]	'per_class_iou': Dict of class_id -> IoU
Dict[str, Any]	'per_class_recall': Dict of class_id -> recall (sensitivity)
Dict[str, Any]	'per_class_precision': Dict of class_id -> precision
Dict[str, Any]	'mean_recall': Mean recall across classes
Dict[str, Any]	'mean_precision': Mean precision across classes

Example

model = torch.load("best_model.pth") results = evaluate_sparse_iou( ... model=model, ... images_dir="tiles/images", ... labels_dir="tiles/labels", ... num_classes=18, ... background_class=0, ... ) print(f"Sparse IoU: {results['mean_sparse_iou']:.4f}")

Source code in geoai/landcover_train.py

def evaluate_sparse_iou(
    model: torch.nn.Module,
    images_dir: str,
    labels_dir: str,
    num_classes: int,
    num_channels: int = 3,
    batch_size: int = 8,
    background_class: int = 0,
    ignore_index: Union[int, bool] = False,
    device: Optional[torch.device] = None,
    verbose: bool = True,
) -> Dict[str, Any]:
    """
    Evaluate a trained model using sparse labels IoU.

    This function is designed for incomplete/sparse ground truth where
    background (0) means "unlabeled" rather than "definitely not this class".
    Predictions in background areas are NOT penalized as false positives.

    Use this for post-training evaluation when your training masks are incomplete.

    Args:
        model: Trained segmentation model
        images_dir: Directory containing validation images
        labels_dir: Directory containing validation labels
        num_classes: Number of classes
        num_channels: Number of input channels (default: 3)
        batch_size: Batch size for evaluation (default: 8)
        background_class: Class ID for background/unlabeled pixels (default: 0)
        ignore_index: Class to ignore during evaluation.
            - If int: specific class index to ignore
            - If False: no class ignored (default)
        device: Torch device (auto-detected if None)
        verbose: Print detailed results (default: True)

    Returns:
        Dictionary containing:
        - 'mean_sparse_iou': Mean IoU across all non-background classes
        - 'per_class_iou': Dict of class_id -> IoU
        - 'per_class_recall': Dict of class_id -> recall (sensitivity)
        - 'per_class_precision': Dict of class_id -> precision
        - 'mean_recall': Mean recall across classes
        - 'mean_precision': Mean precision across classes

    Example:
        >>> model = torch.load("best_model.pth")
        >>> results = evaluate_sparse_iou(
        ...     model=model,
        ...     images_dir="tiles/images",
        ...     labels_dir="tiles/labels",
        ...     num_classes=18,
        ...     background_class=0,
        ... )
        >>> print(f"Sparse IoU: {results['mean_sparse_iou']:.4f}")
    """
    import os
    import rasterio
    from torch.utils.data import DataLoader, Dataset

    if device is None:
        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    model.to(device)
    model.eval()

    # Get all image and label files
    image_extensions = (".tif", ".tiff", ".png", ".jpg", ".jpeg")
    image_files = sorted(
        [
            os.path.join(images_dir, f)
            for f in os.listdir(images_dir)
            if f.lower().endswith(image_extensions)
        ]
    )
    label_files = sorted(
        [
            os.path.join(labels_dir, f)
            for f in os.listdir(labels_dir)
            if f.lower().endswith(image_extensions)
        ]
    )

    if len(image_files) != len(label_files):
        raise ValueError(
            f"Mismatch: {len(image_files)} images vs {len(label_files)} labels"
        )

    if verbose:
        logger.info("=" * 60)
        logger.info("SPARSE LABELS IoU EVALUATION")
        logger.info("=" * 60)
        logger.info("Evaluating %d image-label pairs", len(image_files))
        logger.info(
            "Background class: %d (predictions here NOT penalized)", background_class
        )
        logger.info("Number of classes: %d", num_classes)

    # Accumulate predictions and targets
    all_preds = []
    all_targets = []

    with torch.no_grad():
        for i, (img_path, label_path) in enumerate(zip(image_files, label_files)):
            # Load image
            with rasterio.open(img_path) as src:
                image = src.read()[:num_channels]  # (C, H, W)

            # Load label
            with rasterio.open(label_path) as src:
                label = src.read(1)  # (H, W)

            # Normalize image to 0-1 range
            image = image.astype("float32")
            if image.max() > 1.0:
                image = image / 255.0

            # Convert to tensor and add batch dimension
            image_tensor = torch.from_numpy(image).unsqueeze(0).to(device)
            label_tensor = torch.from_numpy(label).unsqueeze(0)

            # Get prediction
            output = model(image_tensor)
            pred = torch.argmax(output, dim=1).cpu()

            all_preds.append(pred)
            all_targets.append(label_tensor)

            if verbose and (i + 1) % 50 == 0:
                logger.info("   Processed %d/%d tiles...", i + 1, len(image_files))

    # Concatenate all predictions and targets
    all_preds = torch.cat(all_preds, dim=0)
    all_targets = torch.cat(all_targets, dim=0)

    # Calculate sparse IoU
    mean_iou, per_class_ious, recalls, precisions = landcover_iou(
        pred=all_preds,
        target=all_targets,
        num_classes=num_classes,
        ignore_index=ignore_index,
        mode="sparse_labels",
        background_class=background_class,
    )

    # Build results dictionary
    per_class_iou_dict = {}
    per_class_recall_dict = {}
    per_class_precision_dict = {}

    class_idx = 0
    for cls in range(num_classes):
        if cls == background_class:
            continue
        if isinstance(ignore_index, int) and cls == ignore_index:
            continue

        per_class_iou_dict[cls] = per_class_ious[cls]
        per_class_recall_dict[cls] = (
            recalls[class_idx] if class_idx < len(recalls) else 0.0
        )
        per_class_precision_dict[cls] = (
            precisions[class_idx] if class_idx < len(precisions) else 0.0
        )
        class_idx += 1

    # Calculate means (excluding background)
    valid_recalls = [r for r in recalls if r > 0]
    valid_precisions = [p for p in precisions if p > 0]
    mean_recall = sum(valid_recalls) / len(valid_recalls) if valid_recalls else 0.0
    mean_precision = (
        sum(valid_precisions) / len(valid_precisions) if valid_precisions else 0.0
    )

    results = {
        "mean_sparse_iou": mean_iou,
        "per_class_iou": per_class_iou_dict,
        "per_class_recall": per_class_recall_dict,
        "per_class_precision": per_class_precision_dict,
        "mean_recall": mean_recall,
        "mean_precision": mean_precision,
    }

    if verbose:
        logger.info("SPARSE LABELS IoU RESULTS:")
        logger.info(
            "   (Predictions in background areas NOT counted as false positives)"
        )
        logger.info("   %-8s %8s %8s %10s", "Class", "IoU", "Recall", "Precision")
        logger.info("   %s", "-" * 36)
        for cls in sorted(per_class_iou_dict.keys()):
            iou = per_class_iou_dict.get(cls, 0.0)
            recall = per_class_recall_dict.get(cls, 0.0)
            precision = per_class_precision_dict.get(cls, 0.0)
            logger.info("   %-8s %8.4f %8.4f %10.4f", cls, iou, recall, precision)

        logger.info("   %s", "-" * 36)
        logger.info(
            "   %-8s %8.4f %8.4f %10.4f",
            "MEAN",
            mean_iou,
            mean_recall,
            mean_precision,
        )
        logger.info("Sparse IoU evaluation complete!")
        logger.info("=" * 60)

    return results

get_landcover_loss_function(loss_name='crossentropy', num_classes=2, ignore_index=-100, class_weights=None, use_class_weights=False, focal_alpha=1.0, focal_gamma=2.0, device=None, smooth=1.0, tversky_alpha=0.5, tversky_beta=0.5, ufl_lambda=0.5, ufl_gamma=0.75, ufl_delta=0.6, region_weights=None, use_log_cosh=False)

Get loss function configured for landcover segmentation.

Parameters:

Name	Type	Description	Default
loss_name	str	Name of loss function. One of "crossentropy", "focal", "dice", "tversky", "unified_focal" (alias "ufl").	'crossentropy'
num_classes	int	Number of classes.	2
ignore_index	Union[int, bool]	Class index to ignore, or False to not ignore any class.	-100
class_weights	Optional[Tensor]	Manual class weights tensor.	None
use_class_weights	bool	Whether to use class weights.	False
focal_alpha	float	Alpha parameter for focal loss.	1.0
focal_gamma	float	Gamma parameter for focal loss.	2.0
device	Optional[device]	Device to place loss function on.	None
smooth	float	Smoothing constant for Dice / Tversky denominator.	1.0
tversky_alpha	float	False-positive weight for Tversky loss.	0.5
tversky_beta	float	False-negative weight for Tversky loss.	0.5
ufl_lambda	float	Balance between distribution and region components in Unified Focal Loss (1.0 = pure focal CE, 0.0 = pure focal Tversky).	0.5
ufl_gamma	float	Focusing parameter for Unified Focal Loss.	0.75
ufl_delta	float	Tversky false-negative weight for Unified Focal Loss (false-positive weight is 1 - ufl_delta).	0.6
region_weights	Optional[Tensor]	Per-class weights for the region component of Unified Focal Loss. Falls back to class_weights if None.	None
use_log_cosh	bool	Apply log(cosh(loss)) stabilisation in Unified Focal Loss.	False

Returns:

Type	Description
Module	Configured loss function.

Source code in geoai/landcover_train.py

def get_landcover_loss_function(
    loss_name: str = "crossentropy",
    num_classes: int = 2,
    ignore_index: Union[int, bool] = -100,
    class_weights: Optional[torch.Tensor] = None,
    use_class_weights: bool = False,
    focal_alpha: float = 1.0,
    focal_gamma: float = 2.0,
    device: Optional[torch.device] = None,
    smooth: float = 1.0,
    tversky_alpha: float = 0.5,
    tversky_beta: float = 0.5,
    ufl_lambda: float = 0.5,
    ufl_gamma: float = 0.75,
    ufl_delta: float = 0.6,
    region_weights: Optional[torch.Tensor] = None,
    use_log_cosh: bool = False,
) -> nn.Module:
    """Get loss function configured for landcover segmentation.

    Args:
        loss_name: Name of loss function. One of ``"crossentropy"``,
            ``"focal"``, ``"dice"``, ``"tversky"``, ``"unified_focal"``
            (alias ``"ufl"``).
        num_classes: Number of classes.
        ignore_index: Class index to ignore, or ``False`` to not ignore any
            class.
        class_weights: Manual class weights tensor.
        use_class_weights: Whether to use class weights.
        focal_alpha: Alpha parameter for focal loss.
        focal_gamma: Gamma parameter for focal loss.
        device: Device to place loss function on.
        smooth: Smoothing constant for Dice / Tversky denominator.
        tversky_alpha: False-positive weight for Tversky loss.
        tversky_beta: False-negative weight for Tversky loss.
        ufl_lambda: Balance between distribution and region components in
            Unified Focal Loss (``1.0`` = pure focal CE, ``0.0`` = pure
            focal Tversky).
        ufl_gamma: Focusing parameter for Unified Focal Loss.
        ufl_delta: Tversky false-negative weight for Unified Focal Loss
            (false-positive weight is ``1 - ufl_delta``).
        region_weights: Per-class weights for the region component of
            Unified Focal Loss.  Falls back to *class_weights* if ``None``.
        use_log_cosh: Apply ``log(cosh(loss))`` stabilisation in Unified
            Focal Loss.

    Returns:
        Configured loss function.
    """

    if device is None:
        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    loss_name = loss_name.lower()

    # Common helpers
    weights = class_weights if use_class_weights else None
    if weights is not None:
        weights = weights.to(device)
    idx = ignore_index if isinstance(ignore_index, int) else -100

    if loss_name == "crossentropy":
        return LandcoverCrossEntropyLoss(
            weight=weights,
            ignore_index=idx,
            reduction="mean",
        )

    elif loss_name == "focal":
        return FocalLoss(
            alpha=focal_alpha,
            gamma=focal_gamma,
            ignore_index=idx,
            reduction="mean",
            weight=weights,
        )

    elif loss_name == "dice":
        return DiceLoss(
            smooth=smooth,
            ignore_index=idx,
            reduction="mean",
            weight=weights,
        )

    elif loss_name == "tversky":
        return TverskyLoss(
            alpha=tversky_alpha,
            beta=tversky_beta,
            smooth=smooth,
            ignore_index=idx,
            reduction="mean",
            weight=weights,
        )

    elif loss_name in ("unified_focal", "ufl"):
        rw = region_weights
        if rw is not None:
            rw = rw.to(device)
        return UnifiedFocalLoss(
            lambda_=ufl_lambda,
            gamma=ufl_gamma,
            delta=ufl_delta,
            smooth=smooth,
            ignore_index=idx,
            weight=weights,
            region_weight=rw,
            use_log_cosh=use_log_cosh,
        )

    else:
        # Fall back to standard PyTorch loss
        return nn.CrossEntropyLoss(
            weight=weights,
            ignore_index=idx,
            reduction="mean",
        )

landcover_iou(pred, target, num_classes, ignore_index=False, smooth=1e-06, mode='mean', boundary_weight_map=None, background_class=None)

Calculate IoU for landcover classification with multiple weighting options.

Supports four IoU calculation modes: 1. "mean": Simple mean IoU across all classes 2. "perclass_frequency": Weight by per-class pixel frequency 3. "boundary_weighted": Weight by distance to class boundaries 4. "sparse_labels": For incomplete ground truth - only penalize FP where GT is positive (does NOT penalize predictions in unlabeled/background areas)

Parameters:

Name	Type	Description	Default
pred	Tensor	Predicted classes (N, H, W) or logits (N, C, H, W)	required
target	Tensor	Ground truth (N, H, W)	required
num_classes	int	Number of classes	required
ignore_index	Union[int, bool]	Class index to ignore (default: None)	False
smooth	float	Smoothing factor to avoid division by zero	1e-06
mode	str	IoU calculation mode ("mean", "perclass_frequency", "boundary_weighted", "sparse_labels")	'mean'
boundary_weight_map	Optional[Tensor]	Optional boundary weights (N, H, W)	None
background_class	Optional[int]	Background/unlabeled class for sparse_labels mode (default: 0)	None

Returns:

Type	Description
Union[float, Tuple[float, List[float], List[int]]]	If mode == "mean": float (mean IoU)
Union[float, Tuple[float, List[float], List[int]]]	If mode == "perclass_frequency": tuple (weighted IoU, per-class IoUs, class counts)
Union[float, Tuple[float, List[float], List[int]]]	If mode == "boundary_weighted": float (boundary-weighted IoU)
Union[float, Tuple[float, List[float], List[int]]]	If mode == "sparse_labels": tuple (sparse IoU, per-class IoUs, per-class recall, per-class precision)

Source code in geoai/landcover_train.py

def landcover_iou(
    pred: torch.Tensor,
    target: torch.Tensor,
    num_classes: int,
    ignore_index: Union[int, bool] = False,
    smooth: float = 1e-6,
    mode: str = "mean",
    boundary_weight_map: Optional[torch.Tensor] = None,
    background_class: Optional[int] = None,
) -> Union[float, Tuple[float, List[float], List[int]]]:
    """
    Calculate IoU for landcover classification with multiple weighting options.

    Supports four IoU calculation modes:
    1. "mean": Simple mean IoU across all classes
    2. "perclass_frequency": Weight by per-class pixel frequency
    3. "boundary_weighted": Weight by distance to class boundaries
    4. "sparse_labels": For incomplete ground truth - only penalize FP where GT is positive
       (does NOT penalize predictions in unlabeled/background areas)

    Args:
        pred: Predicted classes (N, H, W) or logits (N, C, H, W)
        target: Ground truth (N, H, W)
        num_classes: Number of classes
        ignore_index: Class index to ignore (default: None)
        smooth: Smoothing factor to avoid division by zero
        mode: IoU calculation mode ("mean", "perclass_frequency", "boundary_weighted", "sparse_labels")
        boundary_weight_map: Optional boundary weights (N, H, W)
        background_class: Background/unlabeled class for sparse_labels mode (default: 0)

    Returns:
        If mode == "mean": float (mean IoU)
        If mode == "perclass_frequency": tuple (weighted IoU, per-class IoUs, class counts)
        If mode == "boundary_weighted": float (boundary-weighted IoU)
        If mode == "sparse_labels": tuple (sparse IoU, per-class IoUs, per-class recall, per-class precision)
    """

    # Convert logits to class predictions if needed
    if pred.dim() == 4:
        pred = torch.argmax(pred, dim=1)

    # Ensure correct shape
    if pred.shape != target.shape:
        raise ValueError(f"Shape mismatch: pred {pred.shape}, target {target.shape}")

    # Create mask for valid pixels
    # Handle ignore_index: int means specific class, False means don't ignore
    if isinstance(ignore_index, int):
        valid_mask = target != ignore_index
    else:
        # ignore_index is False or any other non-int value - don't ignore anything
        valid_mask = torch.ones_like(target, dtype=torch.bool)

    # Simple mean IoU
    if mode == "mean":
        ious = []
        for cls in range(num_classes):
            if isinstance(ignore_index, int) and cls == ignore_index:
                continue

            pred_cls = (pred == cls) & valid_mask
            target_cls = (target == cls) & valid_mask

            intersection = (pred_cls & target_cls).sum().float()
            union = (pred_cls | target_cls).sum().float()

            if union > 0:
                iou = (intersection + smooth) / (union + smooth)
                ious.append(iou.item())

        return sum(ious) / len(ious) if ious else 0.0

    # Per-class frequency weighted IoU
    elif mode == "perclass_frequency":
        ious = []
        class_counts = []

        # Filter out ignore_index from target
        if isinstance(ignore_index, int):
            target_filtered = target[valid_mask]
            pred_filtered = pred[valid_mask]
        else:
            target_filtered = target.view(-1)
            pred_filtered = pred.view(-1)

        total_valid_pixels = target_filtered.numel()

        for cls in range(num_classes):
            if isinstance(ignore_index, int) and cls == ignore_index:
                continue

            pred_cls = pred_filtered == cls
            target_cls = target_filtered == cls

            intersection = (pred_cls & target_cls).sum().float()
            union = (pred_cls | target_cls).sum().float()

            class_pixel_count = target_cls.sum().item()

            if union > 0:
                iou = (intersection + smooth) / (union + smooth)
                ious.append(iou.item())
                class_counts.append(class_pixel_count)
            else:
                ious.append(0.0)
                class_counts.append(0)

        # Calculate frequency-weighted IoU
        if sum(class_counts) > 0:
            weights = [count / total_valid_pixels for count in class_counts]
            weighted_iou = sum(iou * weight for iou, weight in zip(ious, weights))
        else:
            weighted_iou = 0.0

        return weighted_iou, ious, class_counts

    # Boundary-weighted IoU
    elif mode == "boundary_weighted":
        if boundary_weight_map is None:
            raise ValueError("boundary_weight_map required for boundary_weighted mode")

        ious = []
        weights = []

        for cls in range(num_classes):
            if isinstance(ignore_index, int) and cls == ignore_index:
                continue

            pred_cls = (pred == cls) & valid_mask
            target_cls = (target == cls) & valid_mask

            # Weight by boundary map
            weighted_intersection = (
                pred_cls & target_cls
            ).float() * boundary_weight_map
            weighted_union = (pred_cls | target_cls).float() * boundary_weight_map

            intersection_sum = weighted_intersection.sum()
            union_sum = weighted_union.sum()

            if union_sum > 0:
                iou = (intersection_sum + smooth) / (union_sum + smooth)
                weight = union_sum.item()
                ious.append(iou.item())
                weights.append(weight)

        if sum(weights) > 0:
            weighted_iou = sum(iou * w for iou, w in zip(ious, weights)) / sum(weights)
        else:
            weighted_iou = 0.0

        return weighted_iou

    # Sparse labels IoU - for incomplete ground truth
    # Key insight: background (0) means "unlabeled", not "definitely not this class"
    # So we DON'T penalize predictions in background areas
    elif mode == "sparse_labels":
        # Default background class is 0 if not specified
        bg_class = background_class if background_class is not None else 0

        ious = []
        recalls = []
        precisions = []
        per_class_ious = []

        # Mask for labeled pixels (ground truth is NOT background)
        labeled_mask = target != bg_class
        if isinstance(ignore_index, int):
            labeled_mask = labeled_mask & (target != ignore_index)

        for cls in range(num_classes):
            # Skip background class and ignore_index
            if cls == bg_class:
                per_class_ious.append(0.0)
                recalls.append(0.0)
                precisions.append(0.0)
                continue
            if isinstance(ignore_index, int) and cls == ignore_index:
                per_class_ious.append(0.0)
                recalls.append(0.0)
                precisions.append(0.0)
                continue

            # Where prediction says this class
            pred_cls = pred == cls
            # Where ground truth says this class
            target_cls = target == cls

            # TRUE POSITIVE: Prediction matches target (both say this class)
            tp = (pred_cls & target_cls).sum().float()

            # FALSE NEGATIVE: Target says this class but prediction doesn't
            fn = (target_cls & ~pred_cls).sum().float()

            # FALSE POSITIVE (SPARSE VERSION):
            # Prediction says this class, target says DIFFERENT class (but NOT background)
            # Key: We don't count predictions in background as FP!
            fp_sparse = (pred_cls & ~target_cls & labeled_mask).sum().float()

            # Standard IoU but with sparse FP definition
            # Union = TP + FN + FP_sparse
            union_sparse = tp + fn + fp_sparse

            if union_sparse > 0:
                iou = (tp + smooth) / (union_sparse + smooth)
                ious.append(iou.item())
                per_class_ious.append(iou.item())
            else:
                per_class_ious.append(0.0)

            # Also compute recall and precision for diagnostic purposes
            # Recall: Of all true positives, how many did we find?
            if (tp + fn) > 0:
                recall = tp / (tp + fn)
                recalls.append(recall.item())
            else:
                recalls.append(0.0)

            # Precision (sparse): Of predictions in labeled areas, how many are correct?
            if (tp + fp_sparse) > 0:
                precision = tp / (tp + fp_sparse)
                precisions.append(precision.item())
            else:
                precisions.append(0.0)

        # Mean IoU across classes (excluding background)
        mean_sparse_iou = sum(ious) / len(ious) if ious else 0.0

        return mean_sparse_iou, per_class_ious, recalls, precisions

    else:
        raise ValueError(
            f"Unknown mode: {mode}. Use 'mean', 'perclass_frequency', 'boundary_weighted', or 'sparse_labels'"
        )

train_segmentation_landcover(images_dir, labels_dir, output_dir, input_format='directory', 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=0.0001, seed=42, val_split=0.2, print_freq=10, verbose=True, save_best_only=True, plot_curves=False, device=None, checkpoint_path=None, resume_training=False, target_size=None, resize_mode='resize', num_workers=None, loss_function='crossentropy', ignore_index=0, use_class_weights=False, focal_alpha=1.0, focal_gamma=2.0, smooth=1.0, tversky_alpha=0.5, tversky_beta=0.5, ufl_lambda=0.5, ufl_gamma=0.75, ufl_delta=0.6, region_weights=None, use_log_cosh=False, custom_multipliers=None, max_class_weight=50.0, use_inverse_frequency=True, validation_iou_mode='standard', boundary_alpha=1.0, background_class=0, training_callback=None, **kwargs)

Train a semantic segmentation model with landcover-specific enhancements.

This is a standalone version that wraps geoai.train.train_segmentation_model with landcover-specific loss functions, class weights, and metrics.

Parameters:

Name	Type	Description	Default
images_dir	str	Directory containing training images.	required
labels_dir	str	Directory containing training labels.	required
output_dir	str	Directory to save model checkpoints and training history.	required
input_format	str	Data format ("directory", "COCO", "YOLO").	'directory'
architecture	str	Model architecture (default: "unet").	'unet'
encoder_name	str	Encoder backbone (default: "resnet34").	'resnet34'
encoder_weights	Optional[str]	Pretrained weights ("imagenet" or None).	'imagenet'
num_channels	int	Number of input channels (default: 3).	3
num_classes	int	Number of output classes (default: 2).	2
batch_size	int	Training batch size (default: 8).	8
num_epochs	int	Number of training epochs (default: 50).	50
learning_rate	float	Initial learning rate (default: 0.001).	0.001
weight_decay	float	Weight decay for optimizer (default: 1e-4).	0.0001
seed	int	Random seed for reproducibility (default: 42).	42
val_split	float	Validation split ratio (default: 0.2).	0.2
print_freq	int	Frequency of training progress prints (default: 10).	10
verbose	bool	Enable verbose output (default: True).	True
save_best_only	bool	Only save best model checkpoint (default: True).	True
plot_curves	bool	Plot training curves at end (default: False).	False
device	Optional[device]	Torch device (auto-detected if None).	None
checkpoint_path	Optional[str]	Path to checkpoint for resuming training.	None
resume_training	bool	Whether to resume from checkpoint (default: False).	False
target_size	Optional[Tuple[int, int]]	Target size for resizing images (H, W) or None.	None
resize_mode	str	How to resize ("resize", "crop", or "pad").	'resize'
num_workers	Optional[int]	Number of dataloader workers (default: auto).	None
loss_function	str	Loss function name. One of "crossentropy", "focal", "dice", "tversky", "unified_focal" (alias "ufl").	'crossentropy'
ignore_index	Union[int, bool]	Class index to ignore during training (default: 0). Set to False so that all pixels contribute to the loss.	0
use_class_weights	bool	Whether to compute and use class weights (default: False).	False
focal_alpha	float	Focal loss alpha parameter (default: 1.0).	1.0
focal_gamma	float	Focal loss gamma parameter (default: 2.0).	2.0
smooth	float	Smoothing constant for Dice / Tversky denominator (default: 1.0).	1.0
tversky_alpha	float	False-positive weight for Tversky loss (default: 0.5).	0.5
tversky_beta	float	False-negative weight for Tversky loss (default: 0.5). Increase relative to tversky_alpha to improve recall.	0.5
ufl_lambda	float	Balance between distribution and region components in Unified Focal Loss (default: 0.5).	0.5
ufl_gamma	float	Focusing parameter for Unified Focal Loss (default: 0.75).	0.75
ufl_delta	float	Tversky false-negative weight inside Unified Focal Loss (default: 0.6).	0.6
region_weights	Optional[Tensor]	Per-class weights for the region component of Unified Focal Loss. Falls back to class_weights if None.	None
use_log_cosh	bool	Apply log(cosh(loss)) stabilisation in Unified Focal Loss (default: False).	False
custom_multipliers	Optional[Dict[int, float]]	Custom class weight multipliers {class_id: multiplier}.	None
max_class_weight	float	Maximum allowed class weight (default: 50.0).	50.0
use_inverse_frequency	bool	Use inverse frequency for weights (default: True).	True
validation_iou_mode	str	IoU calculation mode for validation (default: "standard"). Options: "standard", "perclass_frequency", "boundary_weighted", "sparse_labels".	'standard'
boundary_alpha	float	Boundary importance factor for wIoU mode (default: 1.0).	1.0
background_class	int	Class ID for background/unlabeled pixels in sparse_labels mode (default: 0).	0
training_callback	Optional[callable]	Optional callback function for automatic metric tracking.	None
**kwargs	Any	Additional arguments passed to base training function.	{}

Returns:

Type	Description
Module	Trained model

Example

from landcover_train import train_segmentation_landcover

model = train_segmentation_landcover( ... images_dir="tiles/images", ... labels_dir="tiles/labels", ... output_dir="models/landcover_001", ... num_classes=5, ... loss_function="focal", ... ignore_index=0, # Ignore background ... use_class_weights=True, ... custom_multipliers={1: 1.5, 4: 0.8}, # Boost class 1, reduce class 4 ... max_class_weight=50.0, ... use_inverse_frequency=True, # Use inverse frequency weighting ... validation_iou_mode="boundary_weighted", # Focus on boundaries ... boundary_alpha=2.0, # Moderate boundary emphasis ... )

Source code in geoai/landcover_train.py

def train_segmentation_landcover(
    images_dir: str,
    labels_dir: str,
    output_dir: str,
    input_format: str = "directory",
    architecture: str = "unet",
    encoder_name: str = "resnet34",
    encoder_weights: Optional[str] = "imagenet",
    num_channels: int = 3,
    num_classes: int = 2,
    batch_size: int = 8,
    num_epochs: int = 50,
    learning_rate: float = 0.001,
    weight_decay: float = 1e-4,
    seed: int = 42,
    val_split: float = 0.2,
    print_freq: int = 10,
    verbose: bool = True,
    save_best_only: bool = True,
    plot_curves: bool = False,
    device: Optional[torch.device] = None,
    checkpoint_path: Optional[str] = None,
    resume_training: bool = False,
    target_size: Optional[Tuple[int, int]] = None,
    resize_mode: str = "resize",
    num_workers: Optional[int] = None,
    loss_function: str = "crossentropy",
    ignore_index: Union[int, bool] = 0,
    use_class_weights: bool = False,
    focal_alpha: float = 1.0,
    focal_gamma: float = 2.0,
    smooth: float = 1.0,
    tversky_alpha: float = 0.5,
    tversky_beta: float = 0.5,
    ufl_lambda: float = 0.5,
    ufl_gamma: float = 0.75,
    ufl_delta: float = 0.6,
    region_weights: Optional[torch.Tensor] = None,
    use_log_cosh: bool = False,
    custom_multipliers: Optional[Dict[int, float]] = None,
    max_class_weight: float = 50.0,
    use_inverse_frequency: bool = True,
    validation_iou_mode: str = "standard",
    boundary_alpha: float = 1.0,
    background_class: int = 0,
    training_callback: Optional[callable] = None,
    **kwargs: Any,
) -> torch.nn.Module:
    """Train a semantic segmentation model with landcover-specific enhancements.

    This is a standalone version that wraps geoai.train.train_segmentation_model
    with landcover-specific loss functions, class weights, and metrics.

    Args:
        images_dir: Directory containing training images.
        labels_dir: Directory containing training labels.
        output_dir: Directory to save model checkpoints and training history.
        input_format: Data format (``"directory"``, ``"COCO"``, ``"YOLO"``).
        architecture: Model architecture (default: ``"unet"``).
        encoder_name: Encoder backbone (default: ``"resnet34"``).
        encoder_weights: Pretrained weights (``"imagenet"`` or ``None``).
        num_channels: Number of input channels (default: 3).
        num_classes: Number of output classes (default: 2).
        batch_size: Training batch size (default: 8).
        num_epochs: Number of training epochs (default: 50).
        learning_rate: Initial learning rate (default: 0.001).
        weight_decay: Weight decay for optimizer (default: 1e-4).
        seed: Random seed for reproducibility (default: 42).
        val_split: Validation split ratio (default: 0.2).
        print_freq: Frequency of training progress prints (default: 10).
        verbose: Enable verbose output (default: True).
        save_best_only: Only save best model checkpoint (default: True).
        plot_curves: Plot training curves at end (default: False).
        device: Torch device (auto-detected if ``None``).
        checkpoint_path: Path to checkpoint for resuming training.
        resume_training: Whether to resume from checkpoint (default: False).
        target_size: Target size for resizing images ``(H, W)`` or ``None``.
        resize_mode: How to resize (``"resize"``, ``"crop"``, or ``"pad"``).
        num_workers: Number of dataloader workers (default: auto).
        loss_function: Loss function name. One of ``"crossentropy"``,
            ``"focal"``, ``"dice"``, ``"tversky"``, ``"unified_focal"``
            (alias ``"ufl"``).
        ignore_index: Class index to ignore during training (default: 0).
            Set to ``False`` so that all pixels contribute to the loss.
        use_class_weights: Whether to compute and use class weights
            (default: False).
        focal_alpha: Focal loss alpha parameter (default: 1.0).
        focal_gamma: Focal loss gamma parameter (default: 2.0).
        smooth: Smoothing constant for Dice / Tversky denominator
            (default: 1.0).
        tversky_alpha: False-positive weight for Tversky loss (default: 0.5).
        tversky_beta: False-negative weight for Tversky loss (default: 0.5).
            Increase relative to *tversky_alpha* to improve recall.
        ufl_lambda: Balance between distribution and region components in
            Unified Focal Loss (default: 0.5).
        ufl_gamma: Focusing parameter for Unified Focal Loss (default: 0.75).
        ufl_delta: Tversky false-negative weight inside Unified Focal Loss
            (default: 0.6).
        region_weights: Per-class weights for the region component of
            Unified Focal Loss.  Falls back to *class_weights* if ``None``.
        use_log_cosh: Apply ``log(cosh(loss))`` stabilisation in Unified
            Focal Loss (default: False).
        custom_multipliers: Custom class weight multipliers
            ``{class_id: multiplier}``.
        max_class_weight: Maximum allowed class weight (default: 50.0).
        use_inverse_frequency: Use inverse frequency for weights
            (default: True).
        validation_iou_mode: IoU calculation mode for validation
            (default: ``"standard"``).  Options:
            ``"standard"``, ``"perclass_frequency"``,
            ``"boundary_weighted"``, ``"sparse_labels"``.
        boundary_alpha: Boundary importance factor for wIoU mode
            (default: 1.0).
        background_class: Class ID for background/unlabeled pixels in
            sparse_labels mode (default: 0).
        training_callback: Optional callback function for automatic metric
            tracking.
        **kwargs: Additional arguments passed to base training function.

    Returns:
        Trained model

    Example:
        >>> from landcover_train import train_segmentation_landcover
        >>>
        >>> model = train_segmentation_landcover(
        ...     images_dir="tiles/images",
        ...     labels_dir="tiles/labels",
        ...     output_dir="models/landcover_001",
        ...     num_classes=5,
        ...     loss_function="focal",
        ...     ignore_index=0,  # Ignore background
        ...     use_class_weights=True,
        ...     custom_multipliers={1: 1.5, 4: 0.8},  # Boost class 1, reduce class 4
        ...     max_class_weight=50.0,
        ...     use_inverse_frequency=True,  # Use inverse frequency weighting
        ...     validation_iou_mode="boundary_weighted",  # Focus on boundaries
        ...     boundary_alpha=2.0,  # Moderate boundary emphasis
        ... )
    """

    # Convert ignore_index to format expected by PyTorch loss functions
    if isinstance(ignore_index, bool) and ignore_index is False:
        ignore_idx_for_loss = -100  # PyTorch default (effectively no ignoring)
    elif isinstance(ignore_index, int):
        ignore_idx_for_loss = ignore_index
    else:
        ignore_idx_for_loss = -100

    # Compute class weights if requested
    class_weights_tensor = None
    if use_class_weights:
        if verbose:
            logger.info("=" * 60)
            logger.info("COMPUTING CLASS WEIGHTS")
            logger.info("=" * 60)

        class_weights_tensor = compute_class_weights(
            labels_dir=labels_dir,
            num_classes=num_classes,
            ignore_index=ignore_index,
            custom_multipliers=custom_multipliers,
            max_weight=max_class_weight,
            use_inverse_frequency=use_inverse_frequency,
        )

        if verbose:
            logger.info("=" * 60)

    # Create custom loss function using landcover-specific implementation
    # This ensures ignore_index and class_weights are properly used
    import torch

    device = (
        device
        if device is not None
        else torch.device("cuda" if torch.cuda.is_available() else "cpu")
    )

    if class_weights_tensor is not None:
        class_weights_tensor = class_weights_tensor.to(device)

    # Create the loss function with proper ignore_index support
    criterion = get_landcover_loss_function(
        loss_name=loss_function,
        num_classes=num_classes,
        ignore_index=ignore_idx_for_loss,
        class_weights=class_weights_tensor,
        use_class_weights=use_class_weights,
        focal_alpha=focal_alpha,
        focal_gamma=focal_gamma,
        device=device,
        smooth=smooth,
        tversky_alpha=tversky_alpha,
        tversky_beta=tversky_beta,
        ufl_lambda=ufl_lambda,
        ufl_gamma=ufl_gamma,
        ufl_delta=ufl_delta,
        region_weights=region_weights,
        use_log_cosh=use_log_cosh,
    )

    if verbose:
        logger.info(
            "Created %s loss function with ignore_index=%s",
            loss_function,
            ignore_idx_for_loss,
        )
        if use_class_weights:
            logger.info("Class weights applied: %s", class_weights_tensor)

    # ==========================================================================
    # ALL MODES: Use custom training loop with landcover_iou for model selection
    # ==========================================================================
    # This ensures ALL IoU modes work correctly, not just sparse_labels
    # The base geoai training function ignores validation_iou_mode parameter

    if verbose:
        mode_descriptions = {
            "standard": "STANDARD (unweighted mean IoU)",
            "perclass_frequency": "PER-CLASS FREQUENCY-WEIGHTED IoU",
            "boundary_weighted": f"BOUNDARY-WEIGHTED IoU (wIoU, α={boundary_alpha})",
            "sparse_labels": f"SPARSE LABELS IoU (bg={background_class} ignored)",
        }
        logger.info("=" * 60)
        logger.info(
            "CUSTOM TRAINING LOOP: %s",
            mode_descriptions.get(validation_iou_mode, validation_iou_mode),
        )
        logger.info("=" * 60)
        if validation_iou_mode == "sparse_labels":
            logger.info(
                "Background class: %d (predictions here NOT penalized)",
                background_class,
            )
        elif validation_iou_mode == "boundary_weighted":
            logger.info(
                "Boundary alpha: %s (higher = more focus on boundaries)",
                boundary_alpha,
            )
        elif validation_iou_mode == "perclass_frequency":
            logger.info("Classes weighted by pixel frequency in dataset")
        logger.info(
            "Using %s IoU for model selection during training", validation_iou_mode
        )
        logger.info("=" * 60)

    model = _train_with_custom_iou(
        images_dir=images_dir,
        labels_dir=labels_dir,
        output_dir=output_dir,
        architecture=architecture,
        encoder_name=encoder_name,
        encoder_weights=encoder_weights,
        num_channels=num_channels,
        num_classes=num_classes,
        batch_size=batch_size,
        num_epochs=num_epochs,
        learning_rate=learning_rate,
        weight_decay=weight_decay,
        seed=seed,
        val_split=val_split,
        print_freq=print_freq,
        verbose=verbose,
        save_best_only=save_best_only,
        plot_curves=plot_curves,
        device=device,
        checkpoint_path=checkpoint_path,
        resume_training=resume_training,
        target_size=target_size,
        num_workers=num_workers,
        criterion=criterion,
        validation_iou_mode=validation_iou_mode,
        boundary_alpha=boundary_alpha,
        background_class=background_class,
        ignore_index=(
            ignore_idx_for_loss
            if isinstance(ignore_idx_for_loss, int) and ignore_idx_for_loss != -100
            else False
        ),
        training_callback=training_callback,
        **kwargs,
    )
    return model