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

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

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

    print(f"Computing class weights from {len(label_files)} 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:
            print(f"Warning: Could not read {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:
        print(f"\nApplying custom multipliers: {custom_multipliers}")
        for class_id, multiplier in custom_multipliers.items():
            if class_id < 0 or class_id >= num_classes:
                print(f"Warning: Invalid class_id {class_id}, skipping")
                continue

            original_weight = weights[class_id].item()
            weights[class_id] = weights[class_id] * multiplier
            print(
                f"  Class {class_id}: {original_weight:.4f} × {multiplier} = {weights[class_id].item():.4f}"
            )
    else:
        print("\nNo custom multipliers provided, using computed weights as-is")

    # Apply maximum weight cap to prevent extreme values
    weights_capped = False
    print(f"\nApplying maximum weight cap of {max_weight}...")
    for class_id in range(num_classes):
        if weights[class_id] > max_weight:
            print(
                f"  Class {class_id}: {weights[class_id].item():.4f} → {max_weight} (capped)"
            )
            weights[class_id] = max_weight
            weights_capped = True

    if not weights_capped:
        print("  No weights exceeded the cap")

    print(f"\nClass pixel counts: {dict(class_counts)}")
    print(f"\nFinal 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
        print(
            f"  Class {class_id}: weight={weights[class_id].item():.4f}, "
            f"pixels={pixel_count:,} ({percent:.2f}%)"
        )

    if isinstance(ignore_index, int) and 0 <= ignore_index < num_classes:
        print(f"\nNote: Class {ignore_index} (ignore_index) has weight 0.0")

    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:
        print(f"\n{'='*60}")
        print("SPARSE LABELS IoU EVALUATION")
        print(f"{'='*60}")
        print(f"Evaluating {len(image_files)} image-label pairs")
        print(f"Background class: {background_class} (predictions here NOT penalized)")
        print(f"Number of classes: {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:
                print(f"   Processed {i + 1}/{len(image_files)} tiles...")

    # 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:
        print("\nSPARSE LABELS IoU RESULTS:")
        print(f"   (Predictions in background areas NOT counted as false positives)")
        print(f"\n   {'Class':<8} {'IoU':>8} {'Recall':>8} {'Precision':>10}")
        print(f"   {'-'*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)
            print(f"   {cls:<8} {iou:>8.4f} {recall:>8.4f} {precision:>10.4f}")

        print(f"   {'-'*36}")
        print(
            f"   {'MEAN':<8} {mean_iou:>8.4f} {mean_recall:>8.4f} {mean_precision:>10.4f}"
        )
        print("\nSparse IoU evaluation complete!")
        print(f"{'='*60}\n")

    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)

Get loss function configured for landcover classification.

Parameters:

Name	Type	Description	Default
loss_name	str	Name of loss function ("crossentropy", "focal", "dice", "combo")	'crossentropy'
num_classes	int	Number of classes	2
ignore_index	Union[int, bool]	Class index to ignore, or False to not ignore any class. - If int: pixels with this label value will be ignored during training - If False: no pixels will be ignored (all pixels contribute to loss)	-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

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,
) -> nn.Module:
    """
    Get loss function configured for landcover classification.

    Args:
        loss_name: Name of loss function ("crossentropy", "focal", "dice", "combo")
        num_classes: Number of classes
        ignore_index: Class index to ignore, or False to not ignore any class.
            - If int: pixels with this label value will be ignored during training
            - If False: no pixels will be ignored (all pixels contribute to loss)
        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

    Returns:
        Configured loss function
    """

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

    loss_name = loss_name.lower()

    if loss_name == "crossentropy":
        weights = class_weights if use_class_weights else None
        if weights is not None:
            weights = weights.to(device)

        # Convert ignore_index: int stays as-is, False becomes -100 (PyTorch default)
        idx = ignore_index if isinstance(ignore_index, int) else -100

        return LandcoverCrossEntropyLoss(
            weight=weights,
            ignore_index=idx,
            reduction="mean",
        )

    elif loss_name == "focal":
        weights = class_weights if use_class_weights else None
        if weights is not None:
            weights = weights.to(device)

        # Convert ignore_index: int stays as-is, False becomes -100 (PyTorch default)
        idx = ignore_index if isinstance(ignore_index, int) else -100

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

    else:
        # Fall back to standard PyTorch loss
        weights = class_weights if use_class_weights else None
        if weights is not None:
            weights = weights.to(device)

        # Convert ignore_index: int stays as-is, False becomes -100 (PyTorch default)
        idx = ignore_index if isinstance(ignore_index, int) else -100

        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, 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 ("crossentropy", "focal")	'crossentropy'
ignore_index	Union[int, bool]	Class index to ignore during training. (default: 0) - If int: pixels with this label value will be ignored during training - If False: no pixels will be ignored (all pixels contribute to 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
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") - "standard": Unweighted mean IoU (all classes equal importance) - "perclass_frequency": Frequency-weighted IoU (classes weighted by pixel count) - "boundary_weighted": Boundary-distance weighted IoU (wIoU, focus on edges) - "sparse_labels": For incomplete ground truth - predictions in background areas are NOT penalized. Uses custom training loop with sparse IoU for model selection. BEST FOR INCOMPLETE/SPARSE HABITAT MASKS.	'standard'
boundary_alpha	float	Boundary importance factor for wIoU mode (default: 1.0) Higher values = more focus on boundaries (0.01-100 range)	1.0
background_class	int	Class ID for background/unlabeled pixels in sparse_labels mode (default: 0). Predictions in this class area are NOT counted as false positives.	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,
    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 ("crossentropy", "focal")
        ignore_index: Class index to ignore during training. (default: 0)
            - If int: pixels with this label value will be ignored during training
            - If False: no pixels will be ignored (all pixels contribute to 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)
        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")
            - "standard": Unweighted mean IoU (all classes equal importance)
            - "perclass_frequency": Frequency-weighted IoU (classes weighted by pixel count)
            - "boundary_weighted": Boundary-distance weighted IoU (wIoU, focus on edges)
            - "sparse_labels": For incomplete ground truth - predictions in background
              areas are NOT penalized. Uses custom training loop with sparse IoU for
              model selection. BEST FOR INCOMPLETE/SPARSE HABITAT MASKS.
        boundary_alpha: Boundary importance factor for wIoU mode (default: 1.0)
            Higher values = more focus on boundaries (0.01-100 range)
        background_class: Class ID for background/unlabeled pixels in sparse_labels mode
            (default: 0). Predictions in this class area are NOT counted as false positives.
        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:
            print("\n" + "=" * 60)
            print("COMPUTING CLASS WEIGHTS")
            print("=" * 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:
            print("=" * 60 + "\n")

    # 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,
    )

    if verbose:
        print(
            f"Created {loss_function} loss function with ignore_index={ignore_idx_for_loss}"
        )
        if use_class_weights:
            print(f"Class weights applied: {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)",
        }
        print("\n" + "=" * 60)
        print(
            f"CUSTOM TRAINING LOOP: {mode_descriptions.get(validation_iou_mode, validation_iou_mode)}"
        )
        print("=" * 60)
        if validation_iou_mode == "sparse_labels":
            print(
                f"Background class: {background_class} (predictions here NOT penalized)"
            )
        elif validation_iou_mode == "boundary_weighted":
            print(
                f"Boundary alpha: {boundary_alpha} (higher = more focus on boundaries)"
            )
        elif validation_iou_mode == "perclass_frequency":
            print("Classes weighted by pixel frequency in dataset")
        print(f"Using {validation_iou_mode} IoU for model selection during training")
        print("=" * 60 + "\n")

    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