import glob
import math
import os
import platform
import random
import time
from typing import Any, Callable, Dict, Iterator, List, Optional, Tuple, Union

import matplotlib.pyplot as plt
import numpy as np
import rasterio
import torch
import torch.nn.functional as F
import torch.utils.data
import torchvision
from PIL import Image
from rasterio.windows import Window
from skimage import measure
from sklearn.model_selection import train_test_split
from torch.utils.data import DataLoader, Dataset
from torchvision.models.detection import (
    fasterrcnn_mobilenet_v3_large_fpn,
    fasterrcnn_resnet50_fpn_v2,
    fcos_resnet50_fpn,
    maskrcnn_resnet50_fpn,
    retinanet_resnet50_fpn_v2,
)
from torchvision.models.detection.faster_rcnn import FastRCNNPredictor
from torchvision.models.detection.mask_rcnn import MaskRCNNPredictor
from tqdm import tqdm

from .utils import download_model_from_hf, get_device

# Additional imports for semantic segmentation
try:
    import segmentation_models_pytorch as smp
    from torch.nn import functional as F

    SMP_AVAILABLE = True
except ImportError:
    SMP_AVAILABLE = False

# Additional imports for Lightly Train
try:
    import lightly_train

    LIGHTLY_TRAIN_AVAILABLE = True
except ImportError:
    LIGHTLY_TRAIN_AVAILABLE = False


def parse_coco_annotations(
    coco_json_path: str, images_dir: str, labels_dir: str
) -> Tuple[List[str], List[str]]:
    """
    Parse COCO format annotations and return lists of image and label paths.

    Args:
        coco_json_path (str): Path to COCO annotations JSON file (instances.json).
        images_dir (str): Directory containing image files.
        labels_dir (str): Directory containing label mask files.

    Returns:
        Tuple[List[str], List[str]]: Lists of image paths and corresponding label paths.
    """
    import json

    with open(coco_json_path, "r") as f:
        coco_data = json.load(f)

    # Create mapping from image_id to filename
    image_files = []
    label_files = []

    for img_info in coco_data["images"]:
        img_filename = img_info["file_name"]
        img_path = os.path.join(images_dir, img_filename)

        # Derive label filename (same as image filename)
        label_path = os.path.join(labels_dir, img_filename)

        if os.path.exists(img_path) and os.path.exists(label_path):
            image_files.append(img_path)
            label_files.append(label_path)

    return image_files, label_files


def parse_yolo_annotations(
    data_dir: str, images_subdir: str = "images", labels_subdir: str = "labels"
) -> Tuple[List[str], List[str]]:
    """
    Parse YOLO format annotations and return lists of image and label paths.

    YOLO format structure:
    - data_dir/images/: Contains image files (.tif, .png, .jpg)
    - data_dir/labels/: Contains label masks (.tif, .png) and YOLO .txt files
    - data_dir/classes.txt: Class names (one per line)

    Args:
        data_dir (str): Root directory containing YOLO-format data.
        images_subdir (str): Subdirectory name for images. Defaults to 'images'.
        labels_subdir (str): Subdirectory name for labels. Defaults to 'labels'.

    Returns:
        Tuple[List[str], List[str]]: Lists of image paths and corresponding label paths.
    """
    images_dir = os.path.join(data_dir, images_subdir)
    labels_dir = os.path.join(data_dir, labels_subdir)

    if not os.path.exists(images_dir):
        raise FileNotFoundError(f"Images directory not found: {images_dir}")
    if not os.path.exists(labels_dir):
        raise FileNotFoundError(f"Labels directory not found: {labels_dir}")

    # Get all image files
    image_extensions = (".tif", ".tiff", ".png", ".jpg", ".jpeg")
    image_files = []
    label_files = []

    for img_file in os.listdir(images_dir):
        if img_file.lower().endswith(image_extensions):
            img_path = os.path.join(images_dir, img_file)

            # Find corresponding label mask (same filename)
            label_path = os.path.join(labels_dir, img_file)

            if os.path.exists(label_path):
                image_files.append(img_path)
                label_files.append(label_path)

    return sorted(image_files), sorted(label_files)


DETECTION_MODELS = {
    "fasterrcnn_resnet50_fpn_v2",
    "fasterrcnn_mobilenet_v3_large_fpn",
    "retinanet_resnet50_fpn_v2",
    "fcos_resnet50_fpn",
    "maskrcnn_resnet50_fpn",
}


def model_has_masks(model_name: str) -> bool:
    """Check whether a detection model produces instance masks.

    Args:
        model_name (str): Name of the detection model.

    Returns:
        bool: True if the model outputs instance masks (Mask R-CNN).
    """
    return model_name == "maskrcnn_resnet50_fpn"


def _adjust_backbone_channels(
    model: torch.nn.Module, model_name: str, num_channels: int
) -> None:
    """Modify the backbone's first conv layer to accept a different number of input channels.

    Args:
        model: The detection model.
        model_name: Name used to locate the first conv layer.
        num_channels: Desired number of input channels.
    """
    if num_channels == 3:
        return

    if num_channels < 3:
        raise ValueError(f"num_channels must be at least 3, got {num_channels}")

    # Locate first conv layer based on architecture
    if "mobilenet" in model_name:
        original_layer = model.backbone.body["0"][0]
        parent = model.backbone.body["0"]
        attr = 0
    else:
        # ResNet-based (fasterrcnn_resnet50, retinanet, fcos, maskrcnn)
        original_layer = model.backbone.body.conv1
        parent = model.backbone.body
        attr = "conv1"

    new_layer = torch.nn.Conv2d(
        num_channels,
        original_layer.out_channels,
        kernel_size=original_layer.kernel_size,
        stride=original_layer.stride,
        padding=original_layer.padding,
        bias=original_layer.bias is not None,
    )

    with torch.no_grad():
        min_ch = min(3, num_channels)
        new_layer.weight[:, :min_ch, :, :] = original_layer.weight[:, :min_ch, :, :]
        if num_channels > 3:
            mean_weight = original_layer.weight.mean(dim=1, keepdim=True)
            for i in range(3, num_channels):
                new_layer.weight[:, i : i + 1, :, :] = mean_weight
        if original_layer.bias is not None:
            new_layer.bias = original_layer.bias

    if isinstance(attr, int):
        parent[attr] = new_layer
    else:
        setattr(parent, attr, new_layer)

    # Also update transform normalization stats
    transform = model.transform
    rgb_mean = [0.485, 0.456, 0.406]
    rgb_std = [0.229, 0.224, 0.225]
    mean_of_means = sum(rgb_mean) / len(rgb_mean)
    mean_of_stds = sum(rgb_std) / len(rgb_std)
    transform.image_mean = rgb_mean + [mean_of_means] * (num_channels - 3)
    transform.image_std = rgb_std + [mean_of_stds] * (num_channels - 3)


def get_detection_model(
    model_name: str = "fasterrcnn_resnet50_fpn_v2",
    num_classes: int = 2,
    num_channels: int = 3,
    pretrained: bool = True,
) -> torch.nn.Module:
    """Create a detection model from the supported torchvision architectures.

    Supports Faster R-CNN (ResNet-50 v2, MobileNet v3), RetinaNet, FCOS,
    and Mask R-CNN. The classification head is replaced to match
    ``num_classes`` and the backbone is adjusted for ``num_channels``.

    Args:
        model_name (str): One of the names in :data:`DETECTION_MODELS`.
            Defaults to ``"fasterrcnn_resnet50_fpn_v2"``.
        num_classes (int): Number of output classes including background.
        num_channels (int): Number of input image channels.
        pretrained (bool): Whether to load pretrained weights.

    Returns:
        torch.nn.Module: The detection model with adjusted heads.

    Raises:
        ValueError: If ``model_name`` is not in :data:`DETECTION_MODELS`.
    """
    if model_name not in DETECTION_MODELS:
        raise ValueError(
            f"Unknown model '{model_name}'. " f"Supported: {sorted(DETECTION_MODELS)}"
        )

    # Mask R-CNN: delegate to existing function
    if model_name == "maskrcnn_resnet50_fpn":
        return get_instance_segmentation_model(num_classes, num_channels, pretrained)

    # Build model constructors
    constructors = {
        "fasterrcnn_resnet50_fpn_v2": (
            fasterrcnn_resnet50_fpn_v2,
            "FasterRCNN_ResNet50_FPN_V2_Weights",
        ),
        "fasterrcnn_mobilenet_v3_large_fpn": (
            fasterrcnn_mobilenet_v3_large_fpn,
            "FasterRCNN_MobileNet_V3_Large_FPN_Weights",
        ),
        "retinanet_resnet50_fpn_v2": (
            retinanet_resnet50_fpn_v2,
            "RetinaNet_ResNet50_FPN_V2_Weights",
        ),
        "fcos_resnet50_fpn": (
            fcos_resnet50_fpn,
            "FCOS_ResNet50_FPN_Weights",
        ),
    }

    constructor, weights_name = constructors[model_name]
    weights = (
        getattr(torchvision.models.detection, weights_name).DEFAULT
        if pretrained
        else None
    )
    model = constructor(weights=weights)

    # Replace classification head for the target num_classes
    if model_name.startswith("fasterrcnn"):
        in_features = model.roi_heads.box_predictor.cls_score.in_features
        model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)
    elif model_name in ("retinanet_resnet50_fpn_v2", "fcos_resnet50_fpn"):
        # Replace cls_logits layer preserving spatial kernel and anchor count
        old_cls = model.head.classification_head.cls_logits
        num_anchors = old_cls.out_channels // model.head.classification_head.num_classes
        model.head.classification_head.num_classes = num_classes
        model.head.classification_head.cls_logits = torch.nn.Conv2d(
            old_cls.in_channels,
            num_anchors * num_classes,
            kernel_size=old_cls.kernel_size,
            stride=old_cls.stride,
            padding=old_cls.padding,
        )

    # Adjust backbone for non-RGB input
    _adjust_backbone_channels(model, model_name, num_channels)

    return model


def get_instance_segmentation_model(
    num_classes: int = 2, num_channels: int = 3, pretrained: bool = True
) -> torch.nn.Module:
    """
    Get Mask R-CNN model with custom input channels and output classes.

    Args:
        num_classes (int): Number of output classes (including background).
        num_channels (int): Number of input channels (3 for RGB, 4 for RGBN).
        pretrained (bool): Whether to use pretrained backbone.

    Returns:
        torch.nn.Module: Mask R-CNN model with specified input channels and output classes.

    Raises:
        ValueError: If num_channels is less than 3.
    """
    # Validate num_channels
    if num_channels < 3:
        raise ValueError("num_channels must be at least 3")

    # Load pre-trained model
    model = maskrcnn_resnet50_fpn(
        pretrained=pretrained,
        progress=True,
        weights=(
            torchvision.models.detection.MaskRCNN_ResNet50_FPN_Weights.DEFAULT
            if pretrained
            else None
        ),
    )

    # Modify transform if num_channels is different from 3
    if num_channels != 3:
        # Get the transform
        transform = model.transform

        # Default values are [0.485, 0.456, 0.406] and [0.229, 0.224, 0.225]
        # Calculate means and stds for additional channels
        rgb_mean = [0.485, 0.456, 0.406]
        rgb_std = [0.229, 0.224, 0.225]

        # Extend them to num_channels (use the mean value for additional channels)
        mean_of_means = sum(rgb_mean) / len(rgb_mean)
        mean_of_stds = sum(rgb_std) / len(rgb_std)

        # Create new lists with appropriate length
        transform.image_mean = rgb_mean + [mean_of_means] * (num_channels - 3)
        transform.image_std = rgb_std + [mean_of_stds] * (num_channels - 3)

    # Get number of input features for the classifier
    in_features = model.roi_heads.box_predictor.cls_score.in_features

    # Replace the pre-trained head with a new one
    model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)

    # Get number of input features for mask classifier
    in_features_mask = model.roi_heads.mask_predictor.conv5_mask.in_channels
    hidden_layer = 256

    # Replace mask predictor with a new one
    model.roi_heads.mask_predictor = MaskRCNNPredictor(
        in_features_mask, hidden_layer, num_classes
    )

    # Modify the first layer if num_channels is different from 3
    if num_channels != 3:
        original_layer = model.backbone.body.conv1
        model.backbone.body.conv1 = torch.nn.Conv2d(
            num_channels,
            original_layer.out_channels,
            kernel_size=original_layer.kernel_size,
            stride=original_layer.stride,
            padding=original_layer.padding,
            bias=original_layer.bias is not None,
        )

        # Copy weights from the original 3 channels to the new layer
        with torch.no_grad():
            # Copy the weights for the first 3 channels
            model.backbone.body.conv1.weight[:, :3, :, :] = original_layer.weight

            # Initialize additional channels with the mean of the first 3 channels
            mean_weight = original_layer.weight.mean(dim=1, keepdim=True)
            for i in range(3, num_channels):
                model.backbone.body.conv1.weight[:, i : i + 1, :, :] = mean_weight

            # Copy bias if it exists
            if original_layer.bias is not None:
                model.backbone.body.conv1.bias = original_layer.bias

    return model


class ObjectDetectionDataset(Dataset):
    """Dataset for object detection from GeoTIFF images and labels."""

    def __init__(
        self,
        image_paths: List[str],
        label_paths: List[str],
        transforms: Optional[Callable] = None,
        num_channels: Optional[int] = None,
        instance_labels: bool = False,
    ) -> None:
        """
        Initialize dataset.

        Args:
            image_paths (list): List of paths to image GeoTIFF files.
            label_paths (list): List of paths to label GeoTIFF files.
            transforms (callable, optional): Transformations to apply to images and masks.
            num_channels (int, optional): Number of channels to use from images. If None,
                auto-detected from the first image.
            instance_labels (bool): If True, treat label mask pixel values as
                pre-assigned instance IDs instead of running connected-component
                labeling.  Use this when label masks already encode unique
                integers per object (e.g., Fields of The World dataset).
                Defaults to False.
        """
        self.image_paths = image_paths
        self.label_paths = label_paths
        self.transforms = transforms
        self.instance_labels = instance_labels

        # Auto-detect the number of channels if not specified
        if num_channels is None:
            with rasterio.open(self.image_paths[0]) as src:
                self.num_channels = src.count
        else:
            self.num_channels = num_channels

    def __len__(self) -> int:
        return len(self.image_paths)

    def __getitem__(self, idx: int) -> Tuple[torch.Tensor, Dict[str, torch.Tensor]]:
        # Load image
        with rasterio.open(self.image_paths[idx]) as src:
            # Read as [C, H, W] format
            image = src.read().astype(np.float32)

            # Normalize image to [0, 1] range
            image = image / 255.0

            # Handle different number of channels
            if image.shape[0] > self.num_channels:
                image = image[
                    : self.num_channels
                ]  # Keep only first 4 bands if more exist
            elif image.shape[0] < self.num_channels:
                # Pad with zeros if less than 4 bands
                padded = np.zeros(
                    (self.num_channels, image.shape[1], image.shape[2]),
                    dtype=np.float32,
                )
                padded[: image.shape[0]] = image
                image = padded

            # Convert to CHW tensor
            image = torch.as_tensor(image, dtype=torch.float32)

        # Load label mask
        with rasterio.open(self.label_paths[idx]) as src:
            label_mask = src.read(1)

        if self.instance_labels:
            # Use pre-assigned instance IDs directly (e.g., FTW dataset)
            unique_ids = np.unique(label_mask)
            unique_ids = unique_ids[unique_ids > 0]  # Remove background
            num_instances = len(unique_ids)
            labeled_mask = label_mask
        else:
            # Find instances using connected components on binary mask
            binary_mask = (label_mask > 0).astype(np.uint8)
            labeled_mask, num_instances = measure.label(
                binary_mask, return_num=True, connectivity=2
            )
            unique_ids = range(1, num_instances + 1)

        # Create list to hold masks for each object instance
        masks = []
        boxes = []
        labels = []

        for i in unique_ids:
            # Create mask for this instance
            instance_mask = (labeled_mask == i).astype(np.uint8)

            # Calculate area and filter out tiny instances (noise)
            area = instance_mask.sum()
            if area < 10:  # Minimum area threshold
                continue

            # Find bounding box coordinates
            pos = np.where(instance_mask)
            if len(pos[0]) == 0:  # Skip if mask is empty
                continue

            xmin = np.min(pos[1])
            xmax = np.max(pos[1])
            ymin = np.min(pos[0])
            ymax = np.max(pos[0])

            # Skip invalid boxes
            if xmax <= xmin or ymax <= ymin:
                continue

            # Add small padding to ensure the mask is within the box
            xmin = max(0, xmin - 1)
            ymin = max(0, ymin - 1)
            xmax = min(label_mask.shape[1] - 1, xmax + 1)
            ymax = min(label_mask.shape[0] - 1, ymax + 1)

            boxes.append([xmin, ymin, xmax, ymax])
            masks.append(instance_mask)
            labels.append(1)  # 1 for building class

        # Handle case with no valid instances
        if len(boxes) == 0:
            # Create a dummy target with minimal required fields
            target = {
                "boxes": torch.zeros((0, 4), dtype=torch.float32),
                "labels": torch.zeros((0), dtype=torch.int64),
                "masks": torch.zeros(
                    (0, label_mask.shape[0], label_mask.shape[1]), dtype=torch.uint8
                ),
                "image_id": torch.tensor([idx]),
                "area": torch.zeros((0), dtype=torch.float32),
                "iscrowd": torch.zeros((0), dtype=torch.int64),
            }
        else:
            # Convert to tensors
            boxes = torch.as_tensor(boxes, dtype=torch.float32)
            labels = torch.as_tensor(labels, dtype=torch.int64)
            masks = torch.as_tensor(np.array(masks), dtype=torch.uint8)

            # Calculate area of boxes
            area = (boxes[:, 3] - boxes[:, 1]) * (boxes[:, 2] - boxes[:, 0])

            # Prepare target dictionary
            target = {
                "boxes": boxes,
                "labels": labels,
                "masks": masks,
                "image_id": torch.tensor([idx]),
                "area": area,
                "iscrowd": torch.zeros_like(labels),  # Assume no crowd instances
            }

        # Apply transforms if specified
        if self.transforms is not None:
            image, target = self.transforms(image, target)

        return image, target


class COCODetectionDataset(Dataset):
    """Dataset for multi-class object detection from COCO-format annotations.

    Reads images (JPEG/PNG/GeoTIFF) and loads instance masks and class labels
    from a COCO-format JSON annotation file. Supports multi-class labels
    unlike ObjectDetectionDataset which only produces binary masks.

    Args:
        coco_json_path (str): Path to COCO annotations JSON file.
        images_dir (str): Directory containing image files.
        image_ids (list, optional): Subset of image IDs to use. If None, uses all.
        category_mapping (dict, optional): Mapping from COCO category_id to
            contiguous class labels (1-indexed, 0 is background). If None,
            auto-generated from the annotations file.
        transforms (callable, optional): Transforms to apply to images and targets.
        num_channels (int, optional): Number of image channels. Auto-detected if None.
        min_area (int): Minimum instance area in pixels. Defaults to 10.
    """

    def __init__(
        self,
        coco_json_path: str,
        images_dir: str,
        image_ids: Optional[List[int]] = None,
        category_mapping: Optional[Dict[int, int]] = None,
        transforms: Optional[Callable] = None,
        num_channels: Optional[int] = None,
        min_area: int = 10,
        compute_masks: bool = True,
    ) -> None:
        """Initialize COCO detection dataset.

        Args:
            coco_json_path (str): Path to COCO annotations JSON file.
            images_dir (str): Directory containing image files.
            image_ids (list, optional): Subset of image IDs to use. If None,
                uses all images that have annotations.
            category_mapping (dict, optional): Mapping from COCO category_id to
                contiguous class labels (1-indexed). If None, auto-generated.
            transforms (callable, optional): Transforms to apply.
            num_channels (int, optional): Number of channels. Auto-detected if None.
            min_area (int): Minimum instance area in pixels. Defaults to 10.
            compute_masks (bool): Whether to decode instance segmentation
                masks. Set to False for bbox-only detection models to speed
                up data loading. Defaults to True.
        """
        import json

        self.images_dir = images_dir
        self.transforms = transforms
        self.min_area = min_area
        self.compute_masks = compute_masks

        # Load COCO annotations
        with open(coco_json_path, "r") as f:
            coco_data = json.load(f)

        # Build category mapping
        categories = coco_data.get("categories", [])
        if category_mapping is not None:
            self.category_mapping = category_mapping
        else:
            # Auto-generate contiguous mapping: category_id -> 1, 2, 3, ...
            sorted_cats = sorted(categories, key=lambda c: c["id"])
            self.category_mapping = {
                cat["id"]: idx + 1 for idx, cat in enumerate(sorted_cats)
            }

        self.class_names = ["background"] + [
            cat["name"] for cat in sorted(categories, key=lambda c: c["id"])
        ]
        self.num_classes = len(self.class_names)

        # Build image info lookup
        image_info_map = {img["id"]: img for img in coco_data["images"]}

        # Build annotations grouped by image_id
        annotations_by_image = {}
        for ann in coco_data.get("annotations", []):
            img_id = ann["image_id"]
            if img_id not in annotations_by_image:
                annotations_by_image[img_id] = []
            annotations_by_image[img_id].append(ann)

        # Filter to requested image_ids or all images with annotations.
        # By default, only images that have at least one annotation are
        # included.  Pass explicit image_ids to include negatives.
        if image_ids is not None:
            valid_ids = [img_id for img_id in image_ids if img_id in image_info_map]
        else:
            valid_ids = [
                img_id for img_id in image_info_map if img_id in annotations_by_image
            ]

        # Build final lists, filtering out images whose files don't exist
        self.image_info = []
        self.annotations = []
        for img_id in valid_ids:
            info = image_info_map[img_id]
            img_path = os.path.join(images_dir, info["file_name"])
            if os.path.exists(img_path):
                self.image_info.append(info)
                self.annotations.append(annotations_by_image.get(img_id, []))

        # Auto-detect number of channels from first image
        if num_channels is not None:
            self.num_channels = num_channels
        elif len(self.image_info) > 0:
            first_path = os.path.join(images_dir, self.image_info[0]["file_name"])
            if first_path.lower().endswith((".tif", ".tiff")):
                with rasterio.open(first_path) as src:
                    self.num_channels = src.count
            else:
                img = Image.open(first_path)
                self.num_channels = len(img.getbands())
        else:
            self.num_channels = 3

    def __len__(self) -> int:
        """Return the number of images in the dataset."""
        return len(self.image_info)

    def _polygon_to_mask(
        self, segmentation: List[List[float]], height: int, width: int
    ) -> np.ndarray:
        """Convert COCO polygon segmentation to binary mask.

        Args:
            segmentation (list): List of polygon coordinate lists.
            height (int): Image height.
            width (int): Image width.

        Returns:
            np.ndarray: Binary mask of shape (height, width).
        """
        from skimage.draw import polygon as draw_polygon

        mask = np.zeros((height, width), dtype=np.uint8)
        for poly_coords in segmentation:
            # COCO format: [x1, y1, x2, y2, ..., xn, yn]
            coords = np.array(poly_coords).reshape(-1, 2)
            rr, cc = draw_polygon(coords[:, 1], coords[:, 0], shape=(height, width))
            mask[rr, cc] = 1
        return mask

    def _rle_to_mask(self, rle: Dict, height: int, width: int) -> np.ndarray:
        """Convert COCO RLE segmentation to binary mask.

        Args:
            rle (dict): RLE encoded mask with 'counts' and 'size' keys.
            height (int): Image height.
            width (int): Image width.

        Returns:
            np.ndarray: Binary mask of shape (height, width).
        """
        try:
            from pycocotools import mask as mask_utils

            if isinstance(rle["counts"], list):
                # Uncompressed RLE
                rle_obj = mask_utils.frPyObjects(rle, height, width)
                return mask_utils.decode(rle_obj).astype(np.uint8)
            else:
                return mask_utils.decode(rle).astype(np.uint8)
        except ImportError:
            # Fallback: manual RLE decode
            counts = rle["counts"]
            if isinstance(counts, str):
                # Compressed RLE - need pycocotools for this
                raise ImportError(
                    "pycocotools is required for compressed RLE masks. "
                    "Install with: pip install pycocotools"
                )
            # Uncompressed RLE: alternating 0-runs and 1-runs
            mask = np.zeros(height * width, dtype=np.uint8)
            pos = 0
            for i, count in enumerate(counts):
                if i % 2 == 1:  # Odd indices are foreground runs
                    mask[pos : pos + count] = 1
                pos += count
            return mask.reshape((height, width), order="F")

    def __getitem__(self, idx: int) -> Tuple[torch.Tensor, Dict[str, torch.Tensor]]:
        """Get a single image and its annotations.

        Args:
            idx (int): Index of the image to retrieve.

        Returns:
            Tuple of (image_tensor, target_dict) where target_dict contains
            boxes, labels, masks, image_id, area, and iscrowd fields.
        """
        info = self.image_info[idx]
        img_path = os.path.join(self.images_dir, info["file_name"])
        height = info["height"]
        width = info["width"]

        # Load image
        if img_path.lower().endswith((".tif", ".tiff")):
            with rasterio.open(img_path) as src:
                image = src.read().astype(np.float32)
        else:
            img = Image.open(img_path).convert("RGB")
            image = np.array(img, dtype=np.float32).transpose(2, 0, 1)

        # Normalize to [0, 1]
        if image.max() > 1.0:
            image = image / 255.0

        # Handle channel count
        if image.shape[0] > self.num_channels:
            image = image[: self.num_channels]
        elif image.shape[0] < self.num_channels:
            padded = np.zeros(
                (self.num_channels, image.shape[1], image.shape[2]),
                dtype=np.float32,
            )
            padded[: image.shape[0]] = image
            image = padded

        image = torch.as_tensor(image, dtype=torch.float32)

        # Process annotations
        anns = self.annotations[idx]
        boxes = []
        masks = []
        labels = []

        for ann in anns:
            if self.compute_masks:
                # Decode segmentation mask
                seg = ann.get("segmentation")
                if seg is None:
                    continue

                if isinstance(seg, dict):
                    mask = self._rle_to_mask(seg, height, width)
                elif isinstance(seg, list) and len(seg) > 0:
                    mask = self._polygon_to_mask(seg, height, width)
                else:
                    continue

                # Filter by area
                area_val = mask.sum()
                if area_val < self.min_area:
                    continue

                # Get bounding box from annotation or compute from mask
                if "bbox" in ann:
                    x, y, w, h = ann["bbox"]
                    xmin, ymin = float(x), float(y)
                    xmax, ymax = float(x + w), float(y + h)
                else:
                    pos = np.where(mask)
                    if len(pos[0]) == 0:
                        continue
                    ymin, ymax = int(np.min(pos[0])), int(np.max(pos[0]))
                    xmin, xmax = int(np.min(pos[1])), int(np.max(pos[1]))

                if xmax <= xmin or ymax <= ymin:
                    continue

                boxes.append([xmin, ymin, xmax, ymax])
                masks.append(mask)
                labels.append(self.category_mapping.get(ann["category_id"], 1))
            else:
                # Bbox-only mode: skip expensive mask decoding
                if "bbox" not in ann:
                    continue
                x, y, w, h = ann["bbox"]
                if w < 1 or h < 1:
                    continue
                area_val = ann.get("area", w * h)
                if area_val < self.min_area:
                    continue
                xmin, ymin = float(x), float(y)
                xmax, ymax = float(x + w), float(y + h)
                boxes.append([xmin, ymin, xmax, ymax])
                labels.append(self.category_mapping.get(ann["category_id"], 1))

        # Use the actual COCO image id (not the dataset index)
        coco_image_id = info["id"]

        # Handle case with no valid instances
        if len(boxes) == 0:
            target = {
                "boxes": torch.zeros((0, 4), dtype=torch.float32),
                "labels": torch.zeros((0,), dtype=torch.int64),
                "image_id": torch.tensor([coco_image_id]),
                "area": torch.zeros((0,), dtype=torch.float32),
                "iscrowd": torch.zeros((0,), dtype=torch.int64),
            }
            if self.compute_masks:
                target["masks"] = torch.zeros((0, height, width), dtype=torch.uint8)
        else:
            boxes_t = torch.as_tensor(boxes, dtype=torch.float32)
            labels_t = torch.as_tensor(labels, dtype=torch.int64)
            area = (boxes_t[:, 3] - boxes_t[:, 1]) * (boxes_t[:, 2] - boxes_t[:, 0])
            target = {
                "boxes": boxes_t,
                "labels": labels_t,
                "image_id": torch.tensor([coco_image_id]),
                "area": area,
                "iscrowd": torch.zeros_like(labels_t),
            }
            if self.compute_masks:
                target["masks"] = torch.as_tensor(np.array(masks), dtype=torch.uint8)

        if self.transforms is not None:
            image, target = self.transforms(image, target)

        return image, target


class Compose:
    """Custom compose transform that works with image and target."""

    def __init__(self, transforms: List[Callable]) -> None:
        """
        Initialize compose transform.

        Args:
            transforms (list): List of transforms to apply.
        """
        self.transforms = transforms

    def __call__(
        self, image: torch.Tensor, target: Dict[str, torch.Tensor]
    ) -> Tuple[torch.Tensor, Dict[str, torch.Tensor]]:
        for t in self.transforms:
            image, target = t(image, target)
        return image, target


class ToTensor:
    """Convert numpy.ndarray to tensor."""

    def __call__(
        self, image: torch.Tensor, target: Dict[str, torch.Tensor]
    ) -> Tuple[torch.Tensor, Dict[str, torch.Tensor]]:
        """
        Apply transform to image and target.

        Args:
            image (torch.Tensor): Input image.
            target (dict): Target annotations.

        Returns:
            tuple: Transformed image and target.
        """
        return image, target


class RandomHorizontalFlip:
    """Random horizontal flip transform."""

    def __init__(self, prob: float = 0.5) -> None:
        """
        Initialize random horizontal flip.

        Args:
            prob (float): Probability of applying the flip.
        """
        self.prob = prob

    def __call__(
        self, image: torch.Tensor, target: Dict[str, torch.Tensor]
    ) -> Tuple[torch.Tensor, Dict[str, torch.Tensor]]:
        if random.random() < self.prob:
            # Flip image
            image = torch.flip(image, dims=[2])  # Flip along width dimension

            # Flip masks
            if "masks" in target and len(target["masks"]) > 0:
                target["masks"] = torch.flip(target["masks"], dims=[2])

            # Update boxes
            if "boxes" in target and len(target["boxes"]) > 0:
                boxes = target["boxes"]
                width = image.shape[2]
                boxes[:, 0], boxes[:, 2] = width - boxes[:, 2], width - boxes[:, 0]
                target["boxes"] = boxes

        return image, target


class RandomVerticalFlip:
    """Random vertical flip transform for detection data."""

    def __init__(self, prob: float = 0.5) -> None:
        """Initialize random vertical flip.

        Args:
            prob (float): Probability of applying the flip.
        """
        self.prob = prob

    def __call__(
        self, image: torch.Tensor, target: Dict[str, torch.Tensor]
    ) -> Tuple[torch.Tensor, Dict[str, torch.Tensor]]:
        if random.random() < self.prob:
            image = torch.flip(image, dims=[1])  # Flip along height dimension

            if "masks" in target and len(target["masks"]) > 0:
                target["masks"] = torch.flip(target["masks"], dims=[1])

            if "boxes" in target and len(target["boxes"]) > 0:
                boxes = target["boxes"]
                height = image.shape[1]
                boxes[:, 1], boxes[:, 3] = height - boxes[:, 3], height - boxes[:, 1]
                target["boxes"] = boxes

        return image, target


def get_transform(train: bool) -> torchvision.transforms.Compose:
    """
    Get transforms for data augmentation.

    Args:
        train (bool): Whether to include training-specific transforms.

    Returns:
        Compose: Composed transforms.
    """
    transforms = []
    transforms.append(ToTensor())

    if train:
        transforms.append(RandomHorizontalFlip(0.5))
        transforms.append(RandomVerticalFlip(0.5))

    return Compose(transforms)


def collate_fn(
    batch: List[Tuple[torch.Tensor, Dict[str, torch.Tensor]]],
) -> Tuple[Tuple[torch.Tensor, ...], Tuple[Dict[str, torch.Tensor], ...]]:
    """
    Custom collate function for batching samples.

    Args:
        batch (list): List of (image, target) tuples.

    Returns:
        tuple: Tuple of images and targets.
    """
    return tuple(zip(*batch))


def train_one_epoch(
    model: torch.nn.Module,
    optimizer: torch.optim.Optimizer,
    data_loader: DataLoader,
    device: torch.device,
    epoch: int,
    print_freq: int = 10,
    verbose: bool = True,
) -> float:
    """
    Train the model for one epoch.

    Args:
        model (torch.nn.Module): The model to train.
        optimizer (torch.optim.Optimizer): The optimizer to use.
        data_loader (torch.utils.data.DataLoader): DataLoader for training data.
        device (torch.device): Device to train on.
        epoch (int): Current epoch number.
        print_freq (int): How often to print progress.
        verbose (bool): Whether to print detailed progress.

    Returns:
        float: Average loss for the epoch.
    """
    model.train()
    total_loss = 0
    num_batches = 0

    start_time = time.time()

    for i, (images, targets) in enumerate(data_loader):
        # Move images and targets to device
        images = list(image.to(device) for image in images)
        targets = [{k: v.to(device) for k, v in t.items()} for t in targets]

        # Forward pass
        loss_dict = model(images, targets)
        losses = sum(loss for loss in loss_dict.values())

        # Skip batches with NaN/Inf loss to prevent gradient corruption
        if not torch.isfinite(losses):
            optimizer.zero_grad()
            continue

        # Backward pass
        optimizer.zero_grad()
        losses.backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
        optimizer.step()

        # Track loss
        total_loss += losses.item()
        num_batches += 1

        # Print progress
        if i % print_freq == 0:
            elapsed_time = time.time() - start_time
            if verbose:
                print(
                    f"Epoch: {epoch + 1}, Batch: {i + 1}/{len(data_loader)}, Loss: {losses.item():.4f}, Time: {elapsed_time:.2f}s"
                )
            start_time = time.time()

    # Calculate average loss (only over non-skipped batches)
    avg_loss = total_loss / max(num_batches, 1)
    return avg_loss


def evaluate(
    model: torch.nn.Module,
    data_loader: DataLoader,
    device: torch.device,
    use_mask_iou: bool = True,
) -> Dict[str, float]:
    """
    Evaluate the model on the validation set.

    Args:
        model (torch.nn.Module): The model to evaluate.
        data_loader (torch.utils.data.DataLoader): DataLoader for validation data.
        device (torch.device): Device to evaluate on.
        use_mask_iou (bool): If True, compute IoU from instance masks
            (requires Mask R-CNN). If False, compute IoU from bounding
            boxes. Defaults to True.

    Returns:
        dict: Evaluation metrics including loss and IoU.
    """
    model.eval()

    # Initialize metrics
    total_loss = 0
    num_loss_batches = 0
    iou_scores = []

    with torch.no_grad():
        for images, targets in data_loader:
            # Move to device
            images = list(image.to(device) for image in images)
            targets = [{k: v.to(device) for k, v in t.items()} for t in targets]

            # Mask R-CNN only returns losses in train mode.  We
            # temporarily enable training-mode forward while keeping
            # BatchNorm and Dropout frozen so running stats and
            # stochastic layers are not affected by validation data.
            if len(targets) > 0:
                try:
                    # Freeze BN/Dropout before switching to train mode
                    for mod in model.modules():
                        if isinstance(mod, (torch.nn.BatchNorm2d, torch.nn.Dropout)):
                            mod.eval()
                    model.train()
                    loss_dict = model(images, targets)
                    model.eval()
                    if isinstance(loss_dict, dict):
                        losses = sum(loss for loss in loss_dict.values())
                        if torch.isfinite(losses):
                            total_loss += losses.item()
                            num_loss_batches += 1
                except (RuntimeError, ValueError) as e:
                    import warnings

                    warnings.warn(
                        f"Could not compute val loss for a batch: {e}",
                        stacklevel=2,
                    )
                    model.eval()

            # Get predictions (model must be in eval mode)
            outputs = model(images)

            # Calculate IoU for each image
            for i, output in enumerate(outputs):
                if use_mask_iou:
                    # Mask-based IoU (Mask R-CNN)
                    if (
                        "masks" not in output
                        or len(output["masks"]) == 0
                        or "masks" not in targets[i]
                        or len(targets[i]["masks"]) == 0
                    ):
                        continue

                    pred_masks = (output["masks"].squeeze(1) > 0.5).float()
                    pred_combined = (
                        torch.max(pred_masks, dim=0)[0]
                        if pred_masks.shape[0] > 0
                        else torch.zeros_like(targets[i]["masks"][0])
                    )
                    target_combined = (
                        torch.max(targets[i]["masks"], dim=0)[0]
                        if targets[i]["masks"].shape[0] > 0
                        else torch.zeros_like(pred_combined)
                    )

                    intersection = (pred_combined * target_combined).sum().item()
                    union = ((pred_combined + target_combined) > 0).sum().item()
                    if union > 0:
                        iou_scores.append(intersection / union)
                else:
                    # Box-based IoU (Faster R-CNN, RetinaNet, FCOS)
                    pred_boxes = output["boxes"]
                    target_boxes = targets[i]["boxes"]
                    if len(pred_boxes) == 0 or len(target_boxes) == 0:
                        continue

                    iou_matrix = torchvision.ops.box_iou(pred_boxes, target_boxes)
                    # Best IoU for each ground truth box
                    best_ious = iou_matrix.max(dim=0)[0]
                    iou_scores.append(best_ious.mean().item())

    # Calculate metrics
    avg_loss = total_loss / num_loss_batches if num_loss_batches > 0 else float("inf")
    avg_iou = sum(iou_scores) / len(iou_scores) if iou_scores else 0

    return {"loss": avg_loss, "IoU": avg_iou}


def visualize_predictions(
    model: torch.nn.Module,
    dataset: Dataset,
    device: torch.device,
    num_samples: int = 5,
    output_dir: Optional[str] = None,
) -> None:
    """
    Visualize model predictions.

    Args:
        model (torch.nn.Module): Trained model.
        dataset (torch.utils.data.Dataset): Dataset to visualize.
        device (torch.device): Device to run inference on.
        num_samples (int): Number of samples to visualize.
        output_dir (str, optional): Directory to save visualizations. If None,
            visualizations are displayed but not saved.
    """
    model.eval()

    # Create output directory if needed
    if output_dir:
        os.makedirs(os.path.abspath(output_dir), exist_ok=True)

    # Select random samples
    indices = random.sample(range(len(dataset)), min(num_samples, len(dataset)))

    for idx in indices:
        # Get image and target
        image, target = dataset[idx]

        # Convert to device and add batch dimension
        image = image.to(device)
        image_batch = [image]

        # Get prediction
        with torch.no_grad():
            output = model(image_batch)[0]

        # Convert image from CHW to HWC for display (first 3 bands as RGB)
        rgb_image = image[:3].cpu().numpy()
        rgb_image = np.transpose(rgb_image, (1, 2, 0))
        rgb_image = np.clip(rgb_image, 0, 1)  # Ensure values are in [0,1]

        # Create binary ground truth mask (combine all instances)
        gt_masks = target["masks"].cpu().numpy()
        gt_combined = (
            np.max(gt_masks, axis=0)
            if len(gt_masks) > 0
            else np.zeros((image.shape[1], image.shape[2]), dtype=np.uint8)
        )

        # Create binary prediction mask (combine all instances with score > 0.5)
        pred_masks = output["masks"].cpu().numpy()
        pred_scores = output["scores"].cpu().numpy()
        high_conf_indices = pred_scores > 0.5

        pred_combined = np.zeros((image.shape[1], image.shape[2]), dtype=np.float32)
        if np.any(high_conf_indices):
            for mask in pred_masks[high_conf_indices]:
                # Apply threshold to each predicted mask
                binary_mask = (mask[0] > 0.5).astype(np.float32)
                # Combine with existing masks
                pred_combined = np.maximum(pred_combined, binary_mask)

        # Create figure
        fig, axs = plt.subplots(1, 3, figsize=(15, 5))

        # Show RGB image
        axs[0].imshow(rgb_image)
        axs[0].set_title("RGB Image")
        axs[0].axis("off")

        # Show prediction
        axs[1].imshow(pred_combined, cmap="viridis")
        axs[1].set_title(f"Predicted Buildings: {np.sum(high_conf_indices)} instances")
        axs[1].axis("off")

        # Show ground truth
        axs[2].imshow(gt_combined, cmap="viridis")
        axs[2].set_title(f"Ground Truth: {len(gt_masks)} instances")
        axs[2].axis("off")

        plt.tight_layout()

        # Save or show
        if output_dir:
            plt.savefig(os.path.join(output_dir, f"prediction_{idx}.png"))
            plt.close()
        else:
            plt.show()


def train_MaskRCNN_model(
    images_dir: str,
    labels_dir: str,
    output_dir: str,
    input_format: str = "directory",
    num_channels: int = 3,
    num_classes: int = 2,
    model: Optional[torch.nn.Module] = None,
    pretrained: bool = True,
    pretrained_model_path: Optional[str] = None,
    batch_size: int = 4,
    num_epochs: int = 10,
    learning_rate: float = 0.005,
    seed: int = 42,
    val_split: float = 0.2,
    visualize: bool = False,
    resume_training: bool = False,
    print_freq: int = 10,
    device: Optional[torch.device] = None,
    num_workers: Optional[int] = None,
    verbose: bool = True,
    model_name: str = "maskrcnn_resnet50_fpn",
    instance_labels: bool = False,
) -> torch.nn.Module:
    """Train and evaluate Mask R-CNN model for instance segmentation.

    This function trains a Mask R-CNN model for instance segmentation using the
    provided dataset. It supports loading a pretrained model to either initialize
    the backbone or to continue training from a specific checkpoint.

    Args:
        images_dir (str): Directory containing image GeoTIFF files (for 'directory' format),
            or root directory containing images/ subdirectory (for 'yolo' format),
            or directory containing images (for 'coco' and 'coco_detection' formats).
        labels_dir (str): Directory containing label GeoTIFF files (for 'directory' format),
            or path to COCO annotations JSON file (for 'coco' and 'coco_detection' formats),
            or not used (for 'yolo' format - labels are in images_dir/labels/).
        output_dir (str): Directory to save model checkpoints and results.
        input_format (str): Input data format - 'directory' (default), 'coco',
            'coco_detection', or 'yolo'.
            - 'directory': Standard directory structure with separate images_dir and labels_dir
            - 'coco': COCO JSON format with pre-rendered label TIFFs
            - 'coco_detection': COCO JSON with direct annotation parsing for multi-class
              detection. Automatically sets num_classes from annotations if num_classes is
              2 (default). labels_dir should be the path to the COCO annotations JSON file.
            - 'yolo': YOLO format (images_dir is root with images/ and labels/ subdirectories)
        num_channels (int, optional): Number of input channels. If None, auto-detected.
            Defaults to 3.
        num_classes (int, optional): Number of output classes including background.
            For binary detection (e.g., building vs background), use 2.
            For multi-class detection, use N+1 where N is the number of object
            classes. For 'coco_detection' format, this is auto-detected from
            annotations if left at default (2). Defaults to 2.
        model (torch.nn.Module, optional): Predefined model. If None, a new model is created.
        pretrained (bool): Whether to use pretrained backbone. This is ignored if
            pretrained_model_path is provided. Defaults to True.
        pretrained_model_path (str, optional): Path to a .pth file to load as a
            pretrained model for continued training. Defaults to None.
        batch_size (int): Batch size for training. Defaults to 4.
        num_epochs (int): Number of training epochs. Defaults to 10.
        learning_rate (float): Initial learning rate. Defaults to 0.005.
        seed (int): Random seed for reproducibility. Defaults to 42.
        val_split (float): Fraction of data to use for validation (0-1). Defaults to 0.2.
        visualize (bool): Whether to generate visualizations of model predictions.
            Defaults to False.
        resume_training (bool): If True and pretrained_model_path is provided,
            will try to load optimizer and scheduler states as well. Defaults to False.
        print_freq (int): Frequency of printing training progress. Defaults to 10.
        device (torch.device): Device to train on. If None, uses CUDA if available.
        num_workers (int): Number of workers for data loading. If None, uses 0 on
            macOS and Windows, 8 otherwise.
        verbose (bool): If True, prints detailed training progress. Defaults to True.
        model_name (str): Name of the model architecture. Defaults to
            "maskrcnn_resnet50_fpn".
        instance_labels (bool): If True, treat label mask pixel values as
            pre-assigned instance IDs instead of running connected-component
            labeling. Use this when label masks already encode unique integers
            per object (e.g., Fields of The World dataset). Defaults to False.
    Returns:
        None: Model weights are saved to output_dir.

    Raises:
        FileNotFoundError: If pretrained_model_path is provided but file doesn't exist.
        RuntimeError: If there's an issue loading the pretrained model.
    """

    import datetime

    # Set random seeds for reproducibility
    torch.manual_seed(seed)
    np.random.seed(seed)
    random.seed(seed)
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False

    # Create output directory
    os.makedirs(os.path.abspath(output_dir), exist_ok=True)

    # Get device
    if device is None:
        device = get_device()
    print(f"Using device: {device}")

    # Get all image and label files based on input format
    use_coco_detection = input_format.lower() == "coco_detection"

    if use_coco_detection:
        # COCO detection format: read annotations directly from COCO JSON
        if verbose:
            print(f"Loading COCO detection annotations from {labels_dir}")

        # Create a full dataset to determine image IDs, then split
        _needs_masks = model_has_masks(model_name)
        full_dataset = COCODetectionDataset(
            coco_json_path=labels_dir,
            images_dir=images_dir,
            transforms=None,
            num_channels=num_channels,
            min_area=10,
            compute_masks=_needs_masks,
        )

        # Auto-detect num_classes from annotations when not explicitly set.
        # We check for the default (2) AND that the dataset has more classes,
        # so explicitly passing num_classes=2 for a binary COCO dataset still
        # works (the dataset won't have more than 2 classes in that case).
        detected = full_dataset.num_classes
        if num_classes == 2 and detected > 2:
            num_classes = detected
            if verbose:
                print(
                    f"Auto-detected {num_classes} classes "
                    f"(including background): {full_dataset.class_names}"
                )

        # Split image indices into train and validation
        all_indices = list(range(len(full_dataset)))
        train_indices, val_indices = train_test_split(
            all_indices, test_size=val_split, random_state=seed
        )

        # Get image IDs for each split
        train_image_ids = [full_dataset.image_info[i]["id"] for i in train_indices]
        val_image_ids = [full_dataset.image_info[i]["id"] for i in val_indices]

        print(
            f"Training on {len(train_indices)} images, "
            f"validating on {len(val_indices)} images"
        )

        # Create split datasets
        train_dataset = COCODetectionDataset(
            coco_json_path=labels_dir,
            images_dir=images_dir,
            image_ids=train_image_ids,
            category_mapping=full_dataset.category_mapping,
            transforms=get_transform(train=True),
            num_channels=num_channels,
            min_area=10,
            compute_masks=_needs_masks,
        )
        val_dataset = COCODetectionDataset(
            coco_json_path=labels_dir,
            images_dir=images_dir,
            image_ids=val_image_ids,
            category_mapping=full_dataset.category_mapping,
            transforms=get_transform(train=False),
            num_channels=num_channels,
            min_area=10,
            compute_masks=_needs_masks,
        )

    elif input_format.lower() == "coco":
        # Parse COCO format annotations
        if verbose:
            print(f"Loading COCO format annotations from {labels_dir}")
        # For COCO format, labels_dir is path to instances.json
        # Labels are typically in a "labels" directory parallel to "annotations"
        coco_root = os.path.dirname(os.path.dirname(labels_dir))  # Go up two levels
        labels_directory = os.path.join(coco_root, "labels")
        image_files, label_files = parse_coco_annotations(
            labels_dir, images_dir, labels_directory
        )
    elif input_format.lower() == "yolo":
        # Parse YOLO format annotations
        if verbose:
            print(f"Loading YOLO format data from {images_dir}")
        image_files, label_files = parse_yolo_annotations(images_dir)
    else:
        # Default: directory format
        # Support multiple image formats: GeoTIFF, PNG, JPG, JPEG, TIF, TIFF
        image_extensions = (".tif", ".tiff", ".png", ".jpg", ".jpeg")
        label_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(label_extensions)
            ]
        )

        # Ensure matching files
        if len(image_files) != len(label_files):
            print("Warning: Number of image files and label files don't match!")
            # Find matching files by basename
            basenames = [os.path.basename(f) for f in image_files]
            label_files = [
                os.path.join(labels_dir, os.path.basename(f))
                for f in image_files
                if os.path.exists(os.path.join(labels_dir, os.path.basename(f)))
            ]
            image_files = [
                f
                for f, b in zip(image_files, basenames)
                if os.path.exists(os.path.join(labels_dir, b))
            ]
            print(f"Using {len(image_files)} matching files")

    if not use_coco_detection:
        print(
            f"Found {len(image_files)} image files and {len(label_files)} label files"
        )

        # Split data into train and validation sets
        train_imgs, val_imgs, train_labels, val_labels = train_test_split(
            image_files, label_files, test_size=val_split, random_state=seed
        )

        print(
            f"Training on {len(train_imgs)} images, "
            f"validating on {len(val_imgs)} images"
        )

        # Create datasets
        train_dataset = ObjectDetectionDataset(
            train_imgs,
            train_labels,
            transforms=get_transform(train=True),
            instance_labels=instance_labels,
        )
        val_dataset = ObjectDetectionDataset(
            val_imgs,
            val_labels,
            transforms=get_transform(train=False),
            instance_labels=instance_labels,
        )

    # Create data loaders
    # Use num_workers=0 on macOS and Windows to avoid multiprocessing issues
    # Windows often has issues with multiprocessing in Jupyter notebooks
    # Increase num_workers for better data loading performance
    if num_workers is None:
        num_workers = 0 if platform.system() in ["Darwin", "Windows"] else 8

    train_loader = DataLoader(
        train_dataset,
        batch_size=batch_size,
        shuffle=True,
        collate_fn=collate_fn,
        num_workers=num_workers,
    )

    val_loader = DataLoader(
        val_dataset,
        batch_size=batch_size,
        shuffle=False,
        collate_fn=collate_fn,
        num_workers=num_workers,
    )

    # Initialize model
    if model is None:
        model = get_detection_model(
            model_name=model_name,
            num_classes=num_classes,
            num_channels=num_channels,
            pretrained=pretrained,
        )
    model.to(device)

    # Set up optimizer
    params = [p for p in model.parameters() if p.requires_grad]
    optimizer = torch.optim.SGD(
        params, lr=learning_rate, momentum=0.9, weight_decay=0.0005
    )

    # Set up learning rate scheduler
    lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=5, gamma=0.8)

    # Initialize training variables
    start_epoch = 0
    best_iou = 0

    # Initialize training history
    training_history = {
        "train_loss": [],
        "val_loss": [],
        "val_iou": [],
        "epochs": [],
        "lr": [],
    }

    # Load pretrained model if provided
    if pretrained_model_path:
        if not os.path.exists(pretrained_model_path):
            raise FileNotFoundError(
                f"Pretrained model file not found: {pretrained_model_path}"
            )

        print(f"Loading pretrained model from: {pretrained_model_path}")
        try:
            # Check if it's a full checkpoint or just model weights
            checkpoint = torch.load(pretrained_model_path, map_location=device)

            if isinstance(checkpoint, dict) and "model_state_dict" in checkpoint:
                # It's a checkpoint with extra information
                model.load_state_dict(checkpoint["model_state_dict"])

                if resume_training:
                    # Resume from checkpoint
                    start_epoch = checkpoint.get("epoch", 0) + 1
                    best_iou = checkpoint.get("best_iou", 0)

                    if "optimizer_state_dict" in checkpoint:
                        optimizer.load_state_dict(checkpoint["optimizer_state_dict"])

                    if "scheduler_state_dict" in checkpoint:
                        lr_scheduler.load_state_dict(checkpoint["scheduler_state_dict"])

                    print(f"Resuming training from epoch {start_epoch}")
                    print(f"Previous best IoU: {best_iou:.4f}")
            else:
                # Assume it's just the model weights
                model.load_state_dict(checkpoint)

            print("Pretrained model loaded successfully")
        except Exception as e:
            raise RuntimeError(f"Failed to load pretrained model: {str(e)}")

    # Training loop
    for epoch in range(start_epoch, num_epochs):
        # Train one epoch
        train_loss = train_one_epoch(
            model, optimizer, train_loader, device, epoch, print_freq, verbose
        )

        # Update learning rate
        lr_scheduler.step()

        # Evaluate
        eval_metrics = evaluate(
            model, val_loader, device, use_mask_iou=model_has_masks(model_name)
        )

        # Record training history
        training_history["train_loss"].append(train_loss)
        training_history["val_loss"].append(eval_metrics["loss"])
        training_history["val_iou"].append(eval_metrics["IoU"])
        training_history["epochs"].append(epoch + 1)
        training_history["lr"].append(optimizer.param_groups[0]["lr"])

        # Print metrics
        print(
            f"Epoch {epoch+1}/{num_epochs}: Train Loss: {train_loss:.4f}, Val Loss: {eval_metrics['loss']:.4f}, Val IoU: {eval_metrics['IoU']:.4f}"
        )

        # Save best model
        if eval_metrics["IoU"] > best_iou:
            best_iou = eval_metrics["IoU"]
            print(f"Saving best model with IoU: {best_iou:.4f}")
            torch.save(model.state_dict(), os.path.join(output_dir, "best_model.pth"))

    # Save final model
    torch.save(model.state_dict(), os.path.join(output_dir, "final_model.pth"))

    # Save training history
    torch.save(training_history, os.path.join(output_dir, "training_history.pth"))

    # Load best model for evaluation and visualization
    model.load_state_dict(torch.load(os.path.join(output_dir, "best_model.pth")))

    # Final evaluation
    final_metrics = evaluate(model, val_loader, device)
    print(
        f"Final Evaluation - Loss: {final_metrics['loss']:.4f}, IoU: {final_metrics['IoU']:.4f}"
    )

    # Visualize results
    if visualize:
        print("Generating visualizations...")
        visualize_predictions(
            model,
            val_dataset,
            device,
            num_samples=5,
            output_dir=os.path.join(output_dir, "visualizations"),
        )

    # Save training summary
    with open(os.path.join(output_dir, "training_summary.txt"), "w") as f:
        f.write(
            f"Training completed on: {datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n"
        )
        f.write(f"Total epochs: {num_epochs}\n")
        f.write(f"Best validation IoU: {best_iou:.4f}\n")
        f.write(f"Final validation IoU: {final_metrics['IoU']:.4f}\n")
        f.write(f"Final validation loss: {final_metrics['loss']:.4f}\n")

        if pretrained_model_path:
            f.write(f"Started from pretrained model: {pretrained_model_path}\n")
            if resume_training:
                f.write(f"Resumed training from epoch {start_epoch}\n")

    print(f"Training complete! Trained model saved to {output_dir}")


def inference_on_geotiff(
    model: torch.nn.Module,
    geotiff_path: str,
    output_path: str,
    window_size: int = 512,
    overlap: int = 256,
    confidence_threshold: float = 0.5,
    batch_size: int = 4,
    num_channels: int = 3,
    device: Optional[torch.device] = None,
    **kwargs: Any,
) -> Tuple[str, float]:
    """
    Perform inference on a large GeoTIFF using a sliding window approach with improved blending.

    Args:
        model (torch.nn.Module): Trained model for inference.
        geotiff_path (str): Path to input GeoTIFF file.
        output_path (str): Path to save output mask GeoTIFF.
        window_size (int): Size of sliding window for inference.
        overlap (int): Overlap between adjacent windows.
        confidence_threshold (float): Confidence threshold for predictions (0-1).
        batch_size (int): Batch size for inference.
        num_channels (int): Number of channels to use from the input image.
        device (torch.device, optional): Device to run inference on. If None, uses CUDA if available.
        **kwargs: Additional arguments.

    Returns:
        tuple: Tuple containing output path and inference time in seconds.
    """
    if device is None:
        device = get_device()

    # Put model in evaluation mode
    model.to(device)
    model.eval()

    # Open the GeoTIFF
    with rasterio.open(geotiff_path) as src:
        # Read metadata
        meta = src.meta
        height = src.height
        width = src.width

        # Update metadata for output raster
        out_meta = meta.copy()
        out_meta.update(
            {"count": 1, "dtype": "uint8"}  # Single band for mask  # Binary mask
        )

        # We'll use two arrays:
        # 1. For accumulating predictions
        pred_accumulator = np.zeros((height, width), dtype=np.float32)
        # 2. For tracking how many predictions contribute to each pixel
        count_accumulator = np.zeros((height, width), dtype=np.float32)

        # Calculate the number of windows needed to cover the entire image
        steps_y = math.ceil((height - overlap) / (window_size - overlap))
        steps_x = math.ceil((width - overlap) / (window_size - overlap))

        # Ensure we cover the entire image
        last_y = height - window_size
        last_x = width - window_size

        total_windows = steps_y * steps_x
        print(
            f"Processing {total_windows} windows with size {window_size}x{window_size} and overlap {overlap}..."
        )

        # Create progress bar
        pbar = tqdm(total=total_windows)

        # Process in batches
        batch_inputs = []
        batch_positions = []
        batch_count = 0

        start_time = time.time()

        # Slide window over the image - make sure we cover the entire image
        for i in range(steps_y + 1):  # +1 to ensure we reach the edge
            y = min(i * (window_size - overlap), last_y)
            y = max(0, y)  # Prevent negative indices

            if y > last_y and i > 0:  # Skip if we've already covered the entire height
                continue

            for j in range(steps_x + 1):  # +1 to ensure we reach the edge
                x = min(j * (window_size - overlap), last_x)
                x = max(0, x)  # Prevent negative indices

                if (
                    x > last_x and j > 0
                ):  # Skip if we've already covered the entire width
                    continue

                # Read window
                window = src.read(window=Window(x, y, window_size, window_size))

                # Check if window is valid
                if window.shape[1] != window_size or window.shape[2] != window_size:
                    # This can happen at image edges - adjust window size
                    current_height = window.shape[1]
                    current_width = window.shape[2]
                    if current_height == 0 or current_width == 0:
                        continue  # Skip empty windows
                else:
                    current_height = window_size
                    current_width = window_size

                # Normalize and prepare input
                image = window.astype(np.float32) / 255.0

                # Handle different number of bands
                if image.shape[0] > num_channels:
                    image = image[:num_channels]
                elif image.shape[0] < num_channels:
                    padded = np.zeros(
                        (num_channels, current_height, current_width), dtype=np.float32
                    )
                    padded[: image.shape[0]] = image
                    image = padded

                # Convert to tensor
                image_tensor = torch.tensor(image, device=device)

                # Add to batch
                batch_inputs.append(image_tensor)
                batch_positions.append((y, x, current_height, current_width))
                batch_count += 1

                # Process batch when it reaches the batch size or at the end
                if batch_count == batch_size or (i == steps_y and j == steps_x):
                    # Forward pass
                    with torch.no_grad():
                        outputs = model(batch_inputs)

                    # Process each output in the batch
                    for idx, output in enumerate(outputs):
                        y_pos, x_pos, h, w = batch_positions[idx]

                        # Create weight matrix that gives higher weight to center pixels
                        # This helps with smooth blending at boundaries
                        y_grid, x_grid = np.mgrid[0:h, 0:w]

                        # Calculate distance from each edge
                        dist_from_left = x_grid
                        dist_from_right = w - x_grid - 1
                        dist_from_top = y_grid
                        dist_from_bottom = h - y_grid - 1

                        # Combine distances (minimum distance to any edge)
                        edge_distance = np.minimum.reduce(
                            [
                                dist_from_left,
                                dist_from_right,
                                dist_from_top,
                                dist_from_bottom,
                            ]
                        )

                        # Convert to weight (higher weight for center pixels)
                        # Normalize to [0, 1]
                        edge_distance = np.minimum(edge_distance, overlap / 2)
                        weight = edge_distance / (overlap / 2)

                        # Get masks for predictions above threshold
                        if len(output["scores"]) > 0:
                            # Get all instances that meet confidence threshold
                            keep = output["scores"] > confidence_threshold
                            masks = output["masks"][keep].squeeze(1)

                            # Combine all instances into one mask
                            if len(masks) > 0:
                                combined_mask = torch.max(masks, dim=0)[0] > 0.5
                                combined_mask = (
                                    combined_mask.cpu().numpy().astype(np.float32)
                                )

                                # Apply weight to prediction
                                weighted_pred = combined_mask * weight

                                # Add to accumulators
                                pred_accumulator[
                                    y_pos : y_pos + h, x_pos : x_pos + w
                                ] += weighted_pred
                                count_accumulator[
                                    y_pos : y_pos + h, x_pos : x_pos + w
                                ] += weight

                    # Reset batch
                    batch_inputs = []
                    batch_positions = []
                    batch_count = 0

                    # Update progress bar
                    pbar.update(len(outputs))

        # Close progress bar
        pbar.close()

        # Calculate final mask by dividing accumulated predictions by counts
        # Handle division by zero
        mask = np.zeros((height, width), dtype=np.uint8)
        valid_pixels = count_accumulator > 0
        if np.any(valid_pixels):
            # Average predictions where we have data
            mask[valid_pixels] = (
                pred_accumulator[valid_pixels] / count_accumulator[valid_pixels] > 0.5
            ).astype(np.uint8)

        # Record time
        inference_time = time.time() - start_time
        print(f"Inference completed in {inference_time:.2f} seconds")

        # Save output
        with rasterio.open(output_path, "w", **out_meta) as dst:
            dst.write(mask, 1)

        print(f"Saved prediction to {output_path}")

        return output_path, inference_time


def instance_segmentation_inference_on_geotiff(
    model: torch.nn.Module,
    geotiff_path: str,
    output_path: str,
    window_size: int = 512,
    overlap: int = 256,
    confidence_threshold: float = 0.5,
    nms_threshold: float = 0.3,
    batch_size: int = 4,
    num_channels: int = 3,
    device: Optional[torch.device] = None,
    **kwargs: Any,
) -> Tuple[Dict[str, str], float, List[Dict]]:
    """
    Perform instance segmentation inference on a large GeoTIFF using a
    sliding window approach.

    This function collects all detections first, then applies class-aware
    non-maximum suppression to handle overlapping detections from different
    windows, preventing artifacts. Three single-band rasters are saved:

    - Instance ID raster at ``output_path``.
    - Class label raster at ``output_path`` with a ``_class`` suffix.
    - Confidence score raster at ``output_path`` with a ``_score`` suffix.

    Detections are returned with compact window-sized masks and their
    offsets to avoid materializing full-image arrays per instance.

    Args:
        model (torch.nn.Module): Trained model for inference.
        geotiff_path (str): Path to input GeoTIFF file.
        output_path (str): Path to save the instance ID raster. Class and
            score rasters are saved alongside with ``_class`` and ``_score``
            suffixes.
        window_size (int): Size of sliding window for inference.
        overlap (int): Overlap between adjacent windows.
        confidence_threshold (float): Confidence threshold for predictions
            (0-1).
        nms_threshold (float): IoU threshold for non-maximum suppression.
            Defaults to 0.3.
        batch_size (int): Batch size for inference.
        num_channels (int): Number of channels to use from the input image.
        device (torch.device, optional): Device to run inference on.
            If None, uses CUDA if available.
        **kwargs: Additional arguments.

    Returns:
        tuple: Tuple of (output_paths_dict, inference_time, detections_list).
            output_paths_dict has keys: ``instance``, ``class_label``,
            ``score``. Each detection is a dict with keys: ``mask``
            (np.ndarray, window-sized), ``mask_offset`` (tuple of
            y, x, h, w), ``score`` (float), ``box`` (list in global
            pixel coords), ``label`` (int), ``instance_id`` (int).
    """
    if device is None:
        device = get_device()

    # Put model in evaluation mode
    model.to(device)
    model.eval()

    # Open the GeoTIFF
    with rasterio.open(geotiff_path) as src:
        # Read metadata
        meta = src.meta
        height = src.height
        width = src.width

        # Derive sibling file paths for class and score rasters
        base, ext = os.path.splitext(output_path)
        class_path = f"{base}_class{ext}"
        score_path = f"{base}_score{ext}"

        # Store all detections globally for NMS
        all_detections = []

        # Calculate the number of windows needed to cover the entire image
        steps_y = math.ceil((height - overlap) / (window_size - overlap))
        steps_x = math.ceil((width - overlap) / (window_size - overlap))

        # Ensure we cover the entire image
        last_y = height - window_size
        last_x = width - window_size

        total_windows = steps_y * steps_x
        print(
            f"Processing {total_windows} windows with size {window_size}x{window_size} and overlap {overlap}..."
        )

        # Create progress bar
        pbar = tqdm(total=total_windows)

        # Process in batches
        batch_inputs = []
        batch_positions = []
        batch_count = 0

        start_time = time.time()

        # Slide window over the image
        for i in range(steps_y + 1):  # +1 to ensure we reach the edge
            y = min(i * (window_size - overlap), last_y)
            y = max(0, y)  # Prevent negative indices

            if y > last_y and i > 0:  # Skip if we've already covered the entire height
                continue

            for j in range(steps_x + 1):  # +1 to ensure we reach the edge
                x = min(j * (window_size - overlap), last_x)
                x = max(0, x)  # Prevent negative indices

                if (
                    x > last_x and j > 0
                ):  # Skip if we've already covered the entire width
                    continue

                # Read window
                window = src.read(window=Window(x, y, window_size, window_size))

                # Check if window is valid
                if window.shape[1] == 0 or window.shape[2] == 0:
                    continue

                # Handle edge cases where window might be smaller than expected
                actual_height, actual_width = window.shape[1], window.shape[2]

                # Convert to [C, H, W] format and normalize
                image = window.astype(np.float32) / 255.0

                # Handle different number of channels
                if image.shape[0] > num_channels:
                    image = image[:num_channels]
                elif image.shape[0] < num_channels:
                    # Pad with zeros if less than expected channels
                    padded = np.zeros(
                        (num_channels, image.shape[1], image.shape[2]), dtype=np.float32
                    )
                    padded[: image.shape[0]] = image
                    image = padded

                # Convert to tensor
                image_tensor = torch.tensor(image, device=device)

                # Add to batch
                batch_inputs.append(image_tensor)
                batch_positions.append((y, x, actual_height, actual_width))
                batch_count += 1

                # Process batch when it reaches the batch size or at the end
                if batch_count == batch_size or (i == steps_y and j == steps_x):
                    # Forward pass
                    with torch.no_grad():
                        outputs = model(batch_inputs)

                    # Process each output in the batch
                    for idx, output in enumerate(outputs):
                        y_pos, x_pos, h, w = batch_positions[idx]

                        # Process each detected instance
                        if len(output["scores"]) > 0:
                            # Get instances that meet confidence threshold
                            keep = output["scores"] > confidence_threshold
                            masks = output["masks"][keep].squeeze(1)
                            scores = output["scores"][keep]
                            boxes = output["boxes"][keep]
                            det_labels = output["labels"][keep]

                            # Convert to global coordinates and store
                            for k in range(len(masks)):
                                mask = masks[k].cpu().numpy() > 0.5
                                score = scores[k].cpu().item()
                                box = boxes[k].cpu().numpy()
                                label = det_labels[k].cpu().item()

                                # Convert box to global coordinates
                                global_box = [
                                    box[0] + x_pos,
                                    box[1] + y_pos,
                                    box[2] + x_pos,
                                    box[3] + y_pos,
                                ]

                                # Store compact mask with offset
                                # (avoids full-image arrays before NMS)
                                all_detections.append(
                                    {
                                        "mask": mask,
                                        "mask_offset": (y_pos, x_pos, h, w),
                                        "score": score,
                                        "box": global_box,
                                        "label": label,
                                    }
                                )

                    # Reset batch
                    batch_inputs = []
                    batch_positions = []
                    batch_count = 0

                    # Update progress bar
                    pbar.update(len(outputs))

        # Close progress bar
        pbar.close()

        print(f"Collected {len(all_detections)} detections before NMS")

        # Apply class-aware Non-Maximum Suppression
        if len(all_detections) > 0:
            # Convert to tensors for NMS
            boxes = torch.tensor(
                [det["box"] for det in all_detections], dtype=torch.float32
            )
            scores = torch.tensor(
                [det["score"] for det in all_detections], dtype=torch.float32
            )
            det_labels = torch.tensor(
                [det["label"] for det in all_detections], dtype=torch.int64
            )

            # Apply class-aware NMS
            keep_indices = torchvision.ops.batched_nms(
                boxes, scores, det_labels, nms_threshold
            )

            # Convert to plain int list for reliable Python list indexing
            keep_indices = keep_indices.cpu().tolist()

            # Keep only the selected detections
            final_detections = [all_detections[i] for i in keep_indices]
            print(f"After NMS: {len(final_detections)} detections")

            # Create final output masks
            class_mask = np.zeros((height, width), dtype=np.uint16)
            instance_mask = np.zeros((height, width), dtype=np.uint32)
            score_mask = np.zeros((height, width), dtype=np.float32)

            # Sort by score (highest first) for consistent ordering
            final_detections.sort(key=lambda x: x["score"], reverse=True)

            # Write raster bands from compact masks using window slices
            # (avoids allocating full-image arrays per detection)
            for instance_id, detection in enumerate(final_detections, 1):
                y_pos, x_pos, h, w = detection["mask_offset"]
                detection["instance_id"] = instance_id

                # Crop mask to actual window dimensions (may differ at edges)
                compact_mask = detection["mask"][:h, :w]

                # Write directly into raster arrays using window slice
                region_inst = instance_mask[y_pos : y_pos + h, x_pos : x_pos + w]
                available = (region_inst == 0) & compact_mask
                class_mask[y_pos : y_pos + h, x_pos : x_pos + w][available] = detection[
                    "label"
                ]
                instance_mask[y_pos : y_pos + h, x_pos : x_pos + w][
                    available
                ] = instance_id
                score_mask[y_pos : y_pos + h, x_pos : x_pos + w][available] = detection[
                    "score"
                ]
        else:
            # No detections found
            class_mask = np.zeros((height, width), dtype=np.uint16)
            instance_mask = np.zeros((height, width), dtype=np.uint32)
            score_mask = np.zeros((height, width), dtype=np.float32)
            final_detections = []

        # Record time
        inference_time = time.time() - start_time
        print(f"Instance segmentation completed in {inference_time:.2f} seconds")
        print(f"Final instances: {len(final_detections)}")

        # Save instance raster (uint32)
        inst_meta = meta.copy()
        inst_meta.update({"count": 1, "dtype": "uint32"})
        with rasterio.open(output_path, "w", **inst_meta) as dst:
            dst.write(instance_mask, 1)

        # Save class raster (uint16)
        cls_meta = meta.copy()
        cls_meta.update({"count": 1, "dtype": "uint16"})
        with rasterio.open(class_path, "w", **cls_meta) as dst:
            dst.write(class_mask, 1)

        # Save score raster (float32)
        score_meta = meta.copy()
        score_meta.update({"count": 1, "dtype": "float32"})
        with rasterio.open(score_path, "w", **score_meta) as dst:
            dst.write(score_mask, 1)

        output_paths = {
            "instance": output_path,
            "class_label": class_path,
            "score": score_path,
        }
        print(f"Saved instance raster to {output_path}")
        print(f"Saved class raster to {class_path}")
        print(f"Saved score raster to {score_path}")

        return output_paths, inference_time, final_detections


def multiclass_detection_inference_on_geotiff(
    model: torch.nn.Module,
    geotiff_path: str,
    output_path: str,
    class_names: Optional[List[str]] = None,
    window_size: int = 512,
    overlap: int = 256,
    confidence_threshold: float = 0.5,
    nms_threshold: float = 0.3,
    batch_size: int = 4,
    num_channels: int = 3,
    device: Optional[torch.device] = None,
    **kwargs: Any,
) -> Tuple[str, float, List[Dict]]:
    """Perform multi-class object detection inference on a GeoTIFF.

    Unlike inference_on_geotiff which produces a binary mask, and
    instance_segmentation_inference_on_geotiff which produces instance IDs
    without class info, this function preserves both class labels and
    instance IDs in the output.

    Args:
        model (torch.nn.Module): Trained Mask R-CNN model.
        geotiff_path (str): Path to input GeoTIFF file.
        output_path (str): Path to save output raster (2-band: class + instance).
        class_names (list, optional): List of class names (index 0 = background).
        window_size (int): Size of sliding window for inference. Defaults to 512.
        overlap (int): Overlap between adjacent windows. Defaults to 256.
        confidence_threshold (float): Minimum score threshold. Defaults to 0.5.
        nms_threshold (float): IoU threshold for NMS. Defaults to 0.3.
        batch_size (int): Inference batch size. Defaults to 4.
        num_channels (int): Number of input channels. Defaults to 3.
        device (torch.device, optional): Compute device. If None, auto-detected.
        **kwargs: Additional arguments.

    Returns:
        Tuple of (output_path, inference_time, detections_list) where each
        detection is a dict with keys: mask, score, box, label.
    """
    if device is None:
        device = get_device()

    model.to(device)
    model.eval()

    with rasterio.open(geotiff_path) as src:
        meta = src.meta
        height = src.height
        width = src.width

        out_meta = meta.copy()

        all_detections = []

        stride = window_size - overlap
        # Compute exact number of steps so loop bounds match total_windows
        steps_y = max(1, math.ceil((height - window_size) / stride) + 1)
        steps_x = max(1, math.ceil((width - window_size) / stride) + 1)
        total_windows = steps_y * steps_x

        print(
            f"Processing {total_windows} windows with size "
            f"{window_size}x{window_size} and overlap {overlap}..."
        )

        pbar = tqdm(total=total_windows)
        batch_inputs = []
        batch_positions = []
        batch_count = 0
        start_time = time.time()

        for i in range(steps_y):
            y = min(i * stride, height - window_size)
            y = max(0, y)

            for j in range(steps_x):
                x = min(j * stride, width - window_size)
                x = max(0, x)

                window = src.read(window=Window(x, y, window_size, window_size))
                if window.shape[1] == 0 or window.shape[2] == 0:
                    continue

                actual_height, actual_width = window.shape[1], window.shape[2]
                image = window.astype(np.float32) / 255.0

                if image.shape[0] > num_channels:
                    image = image[:num_channels]
                elif image.shape[0] < num_channels:
                    padded = np.zeros(
                        (num_channels, image.shape[1], image.shape[2]),
                        dtype=np.float32,
                    )
                    padded[: image.shape[0]] = image
                    image = padded

                image_tensor = torch.tensor(image, device=device)
                batch_inputs.append(image_tensor)
                batch_positions.append((y, x, actual_height, actual_width))
                batch_count += 1

                if batch_count == batch_size or (i == steps_y - 1 and j == steps_x - 1):
                    with torch.no_grad():
                        outputs = model(batch_inputs)

                    for idx, output in enumerate(outputs):
                        y_pos, x_pos, h, w = batch_positions[idx]

                        if len(output["scores"]) > 0:
                            keep = output["scores"] > confidence_threshold
                            scores = output["scores"][keep]
                            boxes = output["boxes"][keep]
                            det_labels = output["labels"][keep]
                            has_masks = "masks" in output
                            if has_masks:
                                masks = output["masks"][keep].squeeze(1)

                            for k in range(len(scores)):
                                score = scores[k].cpu().item()
                                box = boxes[k].cpu().numpy()
                                label = det_labels[k].cpu().item()

                                if has_masks:
                                    mask = masks[k].cpu().numpy() > 0.5
                                else:
                                    # Create box-shaped mask for
                                    # bbox-only models
                                    mask = np.zeros((h, w), dtype=bool)
                                    bx = box.astype(int)
                                    mask[
                                        max(0, bx[1]) : min(h, bx[3]),
                                        max(0, bx[0]) : min(w, bx[2]),
                                    ] = True

                                global_box = [
                                    box[0] + x_pos,
                                    box[1] + y_pos,
                                    box[2] + x_pos,
                                    box[3] + y_pos,
                                ]

                                all_detections.append(
                                    {
                                        "mask": mask,
                                        "mask_offset": (y_pos, x_pos, h, w),
                                        "score": score,
                                        "box": global_box,
                                        "label": label,
                                    }
                                )

                    batch_inputs = []
                    batch_positions = []
                    batch_count = 0
                    pbar.update(len(outputs))

        pbar.close()
        print(f"Collected {len(all_detections)} detections before NMS")

        # Apply class-aware Non-Maximum Suppression
        if len(all_detections) > 0:
            boxes = torch.tensor(
                [det["box"] for det in all_detections], dtype=torch.float32
            )
            scores = torch.tensor(
                [det["score"] for det in all_detections], dtype=torch.float32
            )
            det_labels = torch.tensor(
                [det["label"] for det in all_detections], dtype=torch.int64
            )

            keep_indices = torchvision.ops.batched_nms(
                boxes, scores, det_labels, nms_threshold
            )

            final_detections = [all_detections[i] for i in keep_indices]
            print(f"After NMS: {len(final_detections)} detections")

            # Create output rasters (uint8 for class, uint32 for instance)
            class_mask = np.zeros((height, width), dtype=np.uint8)
            instance_mask = np.zeros((height, width), dtype=np.uint32)

            final_detections.sort(key=lambda x: x["score"], reverse=True)

            # Materialize full-size masks only for final kept detections
            for instance_id, detection in enumerate(final_detections, 1):
                compact_mask = detection["mask"]
                y_pos, x_pos, h, w = detection["mask_offset"]
                full_mask = np.zeros((height, width), dtype=bool)
                full_mask[y_pos : y_pos + h, x_pos : x_pos + w] = compact_mask
                # Replace compact mask with full mask for downstream use
                detection["mask"] = full_mask
                del detection["mask_offset"]

                available_pixels = (instance_mask == 0) & full_mask
                class_mask[available_pixels] = detection["label"]
                instance_mask[available_pixels] = instance_id
        else:
            class_mask = np.zeros((height, width), dtype=np.uint8)
            instance_mask = np.zeros((height, width), dtype=np.uint32)
            final_detections = []

        inference_time = time.time() - start_time
        print(f"Multi-class detection completed in {inference_time:.2f} seconds")
        print(f"Final detections: {len(final_detections)}")

        if class_names and len(final_detections) > 0:
            from collections import Counter

            label_counts = Counter(d["label"] for d in final_detections)
            for label_id, count in sorted(label_counts.items()):
                name = (
                    class_names[label_id]
                    if label_id < len(class_names)
                    else f"class_{label_id}"
                )
                print(f"  {name}: {count} detections")

        # Write class band as uint8, instance band as uint32
        out_meta.update({"count": 2, "dtype": "uint32"})
        with rasterio.open(output_path, "w", **out_meta) as dst:
            dst.write(class_mask.astype(np.uint32), 1)
            dst.write(instance_mask, 2)

        print(f"Saved multi-class detection to {output_path}")

        return output_path, inference_time, final_detections


def evaluate_coco_metrics(
    model: torch.nn.Module,
    data_loader: DataLoader,
    device: torch.device,
    class_names: Optional[List[str]] = None,
    iou_thresholds: Optional[List[float]] = None,
    verbose: bool = True,
) -> Dict[str, float]:
    """Evaluate object detection model using COCO-style metrics.

    Computes AP (Average Precision) at various IoU thresholds for each class
    and mAP (mean AP) across all classes.

    Args:
        model (torch.nn.Module): Trained detection model.
        data_loader (DataLoader): Validation data loader.
        device (torch.device): Compute device.
        class_names (list, optional): List of class names (excluding background).
            Index 0 corresponds to class label 1.
        iou_thresholds (list, optional): IoU thresholds for AP calculation.
            Defaults to [0.5, 0.55, ..., 0.95].
        verbose (bool): Whether to print progress. Defaults to True.

    Returns:
        Dict with keys: 'mAP@0.5', 'mAP@0.75', 'mAP@[0.5:0.95]',
        and per-class 'AP@0.5/<class_name>' entries.
    """
    if iou_thresholds is None:
        iou_thresholds = [0.5 + 0.05 * i for i in range(10)]

    model.eval()

    # Collect all predictions and ground truths
    all_predictions = []  # list of dicts per image
    all_targets = []

    if verbose:
        print("Running inference for evaluation...")

    for images, targets in tqdm(data_loader, disable=not verbose):
        images = [img.to(device) for img in images]

        with torch.no_grad():
            outputs = model(images)

        for output, target in zip(outputs, targets):
            pred = {
                "boxes": output["boxes"].cpu(),
                "scores": output["scores"].cpu(),
                "labels": output["labels"].cpu(),
            }
            gt = {
                "boxes": target["boxes"].cpu(),
                "labels": target["labels"].cpu(),
            }
            all_predictions.append(pred)
            all_targets.append(gt)

    # Determine all class IDs present
    all_class_ids = set()
    for gt in all_targets:
        all_class_ids.update(gt["labels"].tolist())
    for pred in all_predictions:
        all_class_ids.update(pred["labels"].tolist())
    all_class_ids = sorted(all_class_ids)

    def _compute_iou_matrix(boxes1, boxes2):
        """Compute IoU between two sets of boxes."""
        x1 = torch.max(boxes1[:, None, 0], boxes2[None, :, 0])
        y1 = torch.max(boxes1[:, None, 1], boxes2[None, :, 1])
        x2 = torch.min(boxes1[:, None, 2], boxes2[None, :, 2])
        y2 = torch.min(boxes1[:, None, 3], boxes2[None, :, 3])

        inter = (x2 - x1).clamp(min=0) * (y2 - y1).clamp(min=0)
        area1 = (boxes1[:, 2] - boxes1[:, 0]) * (boxes1[:, 3] - boxes1[:, 1])
        area2 = (boxes2[:, 2] - boxes2[:, 0]) * (boxes2[:, 3] - boxes2[:, 1])
        union = area1[:, None] + area2[None, :] - inter

        return inter / (union + 1e-6)

    def _compute_ap(precision, recall):
        """Compute AP using 101-point interpolation (COCO style)."""
        recall_points = torch.linspace(0, 1, 101)
        ap = 0.0
        for r in recall_points:
            prec_at_r = precision[recall >= r]
            if len(prec_at_r) > 0:
                ap += prec_at_r.max().item()
        return ap / 101.0

    results = {}

    for iou_thresh in iou_thresholds:
        per_class_ap = []

        for class_id in all_class_ids:
            # Gather all predictions and ground truths for this class
            all_scores = []
            all_tp = []
            total_gt = 0

            for pred, gt in zip(all_predictions, all_targets):
                # Filter predictions for this class
                pred_mask = pred["labels"] == class_id
                pred_boxes = pred["boxes"][pred_mask]
                pred_scores = pred["scores"][pred_mask]

                # Filter ground truths for this class
                gt_mask = gt["labels"] == class_id
                gt_boxes = gt["boxes"][gt_mask]
                total_gt += len(gt_boxes)

                if len(pred_boxes) == 0:
                    continue

                # Sort by score (descending)
                sort_idx = torch.argsort(pred_scores, descending=True)
                pred_boxes = pred_boxes[sort_idx]
                pred_scores = pred_scores[sort_idx]

                matched_gt = set()

                for pi in range(len(pred_boxes)):
                    all_scores.append(pred_scores[pi].item())

                    if len(gt_boxes) == 0:
                        all_tp.append(0)
                        continue

                    iou_matrix = _compute_iou_matrix(pred_boxes[pi : pi + 1], gt_boxes)
                    max_iou, max_idx = iou_matrix[0].max(dim=0)

                    if (
                        max_iou.item() >= iou_thresh
                        and max_idx.item() not in matched_gt
                    ):
                        all_tp.append(1)
                        matched_gt.add(max_idx.item())
                    else:
                        all_tp.append(0)

            if total_gt == 0:
                continue

            # Sort by score
            if len(all_scores) == 0:
                per_class_ap.append(0.0)
                continue

            sort_idx = sorted(
                range(len(all_scores)),
                key=lambda k: all_scores[k],
                reverse=True,
            )
            tp_sorted = torch.tensor([all_tp[i] for i in sort_idx])

            tp_cumsum = torch.cumsum(tp_sorted, dim=0).float()
            fp_cumsum = torch.cumsum(1 - tp_sorted, dim=0).float()

            precision = tp_cumsum / (tp_cumsum + fp_cumsum)
            recall = tp_cumsum / total_gt

            ap = _compute_ap(precision, recall)
            per_class_ap.append(ap)

            # Store per-class AP at 0.5
            if abs(iou_thresh - 0.5) < 1e-6 and class_names is not None:
                idx = class_id - 1  # class_id is 1-indexed
                name = (
                    class_names[idx] if idx < len(class_names) else f"class_{class_id}"
                )
                results[f"AP@0.5/{name}"] = ap

        mean_ap = sum(per_class_ap) / len(per_class_ap) if per_class_ap else 0.0

        if abs(iou_thresh - 0.5) < 1e-6:
            results["mAP@0.5"] = mean_ap
        if abs(iou_thresh - 0.75) < 1e-6:
            results["mAP@0.75"] = mean_ap

    # Compute mAP@[0.5:0.95]
    all_maps = []
    for iou_thresh in iou_thresholds:
        per_class_ap = []
        for class_id in all_class_ids:
            all_scores = []
            all_tp = []
            total_gt = 0

            for pred, gt in zip(all_predictions, all_targets):
                pred_mask = pred["labels"] == class_id
                pred_boxes = pred["boxes"][pred_mask]
                pred_scores = pred["scores"][pred_mask]
                gt_mask = gt["labels"] == class_id
                gt_boxes = gt["boxes"][gt_mask]
                total_gt += len(gt_boxes)

                if len(pred_boxes) == 0:
                    continue

                sort_idx = torch.argsort(pred_scores, descending=True)
                pred_boxes = pred_boxes[sort_idx]
                pred_scores = pred_scores[sort_idx]
                matched_gt = set()

                for pi in range(len(pred_boxes)):
                    all_scores.append(pred_scores[pi].item())
                    if len(gt_boxes) == 0:
                        all_tp.append(0)
                        continue
                    iou_matrix = _compute_iou_matrix(pred_boxes[pi : pi + 1], gt_boxes)
                    max_iou, max_idx = iou_matrix[0].max(dim=0)
                    if (
                        max_iou.item() >= iou_thresh
                        and max_idx.item() not in matched_gt
                    ):
                        all_tp.append(1)
                        matched_gt.add(max_idx.item())
                    else:
                        all_tp.append(0)

            if total_gt == 0 or len(all_scores) == 0:
                per_class_ap.append(0.0)
                continue

            sort_idx = sorted(
                range(len(all_scores)),
                key=lambda k: all_scores[k],
                reverse=True,
            )
            tp_sorted = torch.tensor([all_tp[i] for i in sort_idx])
            tp_cumsum = torch.cumsum(tp_sorted, dim=0).float()
            fp_cumsum = torch.cumsum(1 - tp_sorted, dim=0).float()
            precision = tp_cumsum / (tp_cumsum + fp_cumsum)
            recall = tp_cumsum / total_gt
            ap = _compute_ap(precision, recall)
            per_class_ap.append(ap)

        mean_ap = sum(per_class_ap) / len(per_class_ap) if per_class_ap else 0.0
        all_maps.append(mean_ap)

    results["mAP@[0.5:0.95]"] = sum(all_maps) / len(all_maps) if all_maps else 0.0

    if verbose:
        print(f"\nEvaluation Results:")
        print(f"  mAP@0.5:        {results.get('mAP@0.5', 0.0):.4f}")
        print(f"  mAP@0.75:       {results.get('mAP@0.75', 0.0):.4f}")
        print(f"  mAP@[0.5:0.95]: {results['mAP@[0.5:0.95]']:.4f}")
        if class_names:
            print(f"\n  Per-class AP@0.5:")
            for key, value in sorted(results.items()):
                if key.startswith("AP@0.5/"):
                    print(f"    {key}: {value:.4f}")

    return results


def object_detection(
    input_path: str,
    output_path: str,
    model_path: str,
    window_size: int = 512,
    overlap: int = 256,
    confidence_threshold: float = 0.5,
    batch_size: int = 4,
    num_channels: int = 3,
    num_classes: int = 2,
    model: Optional[torch.nn.Module] = None,
    pretrained: bool = True,
    device: Optional[torch.device] = None,
    **kwargs: Any,
) -> None:
    """
    Perform object detection on a GeoTIFF using a pre-trained Mask R-CNN model.

    Args:
        input_path (str): Path to input GeoTIFF file.
        output_path (str): Path to save output mask GeoTIFF.
        model_path (str): Path to trained model weights.
        window_size (int): Size of sliding window for inference.
        overlap (int): Overlap between adjacent windows.
        confidence_threshold (float): Confidence threshold for predictions (0-1).
        batch_size (int): Batch size for inference.
        num_channels (int): Number of channels in the input image and model.
        num_classes (int): Number of output classes including background. Defaults to 2.
        model (torch.nn.Module, optional): Predefined model. If None, a new model is created.
        pretrained (bool): Whether to use pretrained backbone for model loading.
        device (torch.device, optional): Device to run inference on. If None, uses CUDA if available.
        **kwargs: Additional arguments passed to inference_on_geotiff.

    Returns:
        None: Output mask is saved to output_path.
    """
    # Load your trained model
    if device is None:
        device = get_device()
    if model is None:
        model = get_instance_segmentation_model(
            num_classes=num_classes, num_channels=num_channels, pretrained=pretrained
        )

    if not os.path.exists(model_path):
        try:
            model_path = download_model_from_hf(model_path)
        except Exception as e:
            raise FileNotFoundError(f"Model file not found: {model_path}")

    # Load state dict and handle DataParallel module prefix
    state_dict = torch.load(model_path, map_location=device)

    # Remove 'module.' prefix if present (from DataParallel training)
    if any(key.startswith("module.") for key in state_dict.keys()):
        state_dict = {
            key.replace("module.", ""): value for key, value in state_dict.items()
        }

    model.load_state_dict(state_dict)
    model.to(device)
    model.eval()

    inference_on_geotiff(
        model=model,
        geotiff_path=input_path,
        output_path=output_path,
        window_size=window_size,  # Adjust based on your model and memory
        overlap=overlap,  # Overlap to avoid edge artifacts
        confidence_threshold=confidence_threshold,
        batch_size=batch_size,  # Adjust based on your GPU memory
        num_channels=num_channels,
        device=device,
        **kwargs,
    )


def object_detection_batch(
    input_paths: Union[str, List[str]],
    output_dir: str,
    model_path: str,
    filenames: Optional[List[str]] = None,
    window_size: int = 512,
    overlap: int = 256,
    confidence_threshold: float = 0.5,
    batch_size: int = 4,
    model: Optional[torch.nn.Module] = None,
    num_channels: int = 3,
    num_classes: int = 2,
    pretrained: bool = True,
    device: Optional[torch.device] = None,
    **kwargs: Any,
) -> None:
    """
    Perform object detection on a GeoTIFF using a pre-trained Mask R-CNN model.

    Args:
        input_paths (str or list): Path(s) to input GeoTIFF file(s). Can be a
            single file path (.tif or .tiff, case-insensitive), a directory
            path (all .tif and .tiff files will be processed), or a list of
            file paths.
        output_dir (str): Directory to save output mask GeoTIFF files.
        model_path (str): Path to trained model weights.
        filenames (list, optional): List of output filenames. If None, defaults to
            "<input_filename>_mask.tif" for each input file.
            If provided, must match the number of input files.
        window_size (int): Size of sliding window for inference.
        overlap (int): Overlap between adjacent windows.
        confidence_threshold (float): Confidence threshold for predictions (0-1).
        batch_size (int): Batch size for inference.
        num_channels (int): Number of channels in the input image and model.
        num_classes (int): Number of output classes including background. Defaults to 2.
        model (torch.nn.Module, optional): Predefined model. If None, a new model is created.
        pretrained (bool): Whether to use pretrained backbone for model loading.
        device (torch.device, optional): Device to run inference on. If None, uses CUDA if available.
        **kwargs: Additional arguments passed to inference_on_geotiff.

    Returns:
        None: Output mask is saved to output_path.
    """
    # Load your trained model
    if device is None:
        device = get_device()

    if model is None:
        model = get_instance_segmentation_model(
            num_classes=num_classes, num_channels=num_channels, pretrained=pretrained
        )

    if not os.path.exists(output_dir):
        os.makedirs(os.path.abspath(output_dir), exist_ok=True)

    if not os.path.exists(model_path):
        try:
            model_path = download_model_from_hf(model_path)
        except Exception as e:
            raise FileNotFoundError(f"Model file not found: {model_path}")

    # Load state dict and handle DataParallel module prefix
    state_dict = torch.load(model_path, map_location=device)

    # Remove 'module.' prefix if present (from DataParallel training)
    if any(key.startswith("module.") for key in state_dict.keys()):
        state_dict = {
            key.replace("module.", ""): value for key, value in state_dict.items()
        }

    model.load_state_dict(state_dict)
    model.to(device)
    model.eval()

    if isinstance(input_paths, list):
        files = input_paths
    elif isinstance(input_paths, str) and not input_paths.lower().endswith(
        (".tif", ".tiff")
    ):
        files = sorted(
            glob.glob(os.path.join(input_paths, "*.tif"))
            + glob.glob(os.path.join(input_paths, "*.tiff"))
        )
    else:
        files = [input_paths]

    if filenames is None:
        filenames = [
            os.path.join(
                output_dir,
                os.path.splitext(os.path.basename(f))[0] + "_mask.tif",
            )
            for f in files
        ]
    else:
        if len(filenames) != len(files):
            raise ValueError("Number of filenames must match number of input files.")

    for index, file in enumerate(files):
        print(f"Processing file {index + 1}/{len(files)}: {file}")
        inference_on_geotiff(
            model=model,
            geotiff_path=file,
            output_path=filenames[index],
            window_size=window_size,  # Adjust based on your model and memory
            overlap=overlap,  # Overlap to avoid edge artifacts
            confidence_threshold=confidence_threshold,
            batch_size=batch_size,  # Adjust based on your GPU memory
            num_channels=num_channels,
            device=device,
            **kwargs,
        )


class SemanticSegmentationDataset(Dataset):
    """Dataset for semantic segmentation from GeoTIFF, PNG, JPG, and other image formats."""

    def __init__(
        self,
        image_paths: List[str],
        label_paths: List[str],
        transforms: Optional[Callable] = None,
        num_channels: Optional[int] = None,
        target_size: Optional[Tuple[int, int]] = None,
        resize_mode: str = "resize",
        num_classes: int = 2,
    ) -> None:
        """
        Initialize dataset for semantic segmentation.

        Args:
            image_paths (list): List of paths to image files (GeoTIFF, PNG, JPG, etc.).
            label_paths (list): List of paths to label files (GeoTIFF, PNG, JPG, etc.).
            transforms (callable, optional): Transformations to apply to images and masks.
            num_channels (int, optional): Number of channels to use from images. If None,
                auto-detected from the first image.
            target_size (tuple, optional): Target size (height, width) for standardizing images.
                If None, images will keep their original sizes.
            resize_mode (str): How to handle size standardization. Options:
                'resize' - Resize images to target_size (may change aspect ratio)
                'pad' - Pad images to target_size (preserves aspect ratio)
            num_classes (int): Number of classes for segmentation. Used for mask normalization.
        """
        self.image_paths = image_paths
        self.label_paths = label_paths
        self.transforms = transforms
        self.target_size = target_size
        self.resize_mode = resize_mode
        self.num_classes = num_classes

        # Auto-detect the number of channels if not specified
        if num_channels is None:
            self.num_channels = self._get_num_channels(self.image_paths[0])
        else:
            self.num_channels = num_channels

    def _is_geotiff(self, file_path: str) -> bool:
        """Check if file is a GeoTIFF based on extension."""
        return file_path.lower().endswith((".tif", ".tiff"))

    def _get_num_channels(self, image_path: str) -> int:
        """Get number of channels from an image file."""
        if self._is_geotiff(image_path):
            with rasterio.open(image_path) as src:
                return src.count
        else:
            # For standard image formats, use PIL
            with Image.open(image_path) as img:
                if img.mode == "RGB":
                    return 3
                elif img.mode == "RGBA":
                    return 4
                elif img.mode == "L":
                    return 1
                else:
                    # Convert to RGB and return 3 channels
                    return 3

    def _resize_image_and_mask(
        self, image: np.ndarray, mask: np.ndarray
    ) -> Tuple[np.ndarray, np.ndarray]:
        """Resize image and mask to target size."""
        if self.target_size is None:
            return image, mask

        target_h, target_w = self.target_size

        if self.resize_mode == "resize":
            # Direct resize (may change aspect ratio)
            image = F.interpolate(
                image.unsqueeze(0),
                size=(target_h, target_w),
                mode="bilinear",
                align_corners=False,
            ).squeeze(0)

            mask = (
                F.interpolate(
                    mask.unsqueeze(0).unsqueeze(0).float(),
                    size=(target_h, target_w),
                    mode="nearest",
                )
                .squeeze(0)
                .squeeze(0)
                .long()
            )
            # Clamp mask values to ensure they're within valid range [0, num_classes-1]
            mask = torch.clamp(mask, 0, self.num_classes - 1)

        elif self.resize_mode == "pad":
            # Pad to target size (preserves aspect ratio)
            image = self._pad_to_size(image, (target_h, target_w))
            mask = self._pad_to_size(mask.unsqueeze(0), (target_h, target_w)).squeeze(0)
            # Clamp mask values to ensure they're within valid range [0, num_classes-1]
            mask = torch.clamp(mask, 0, self.num_classes - 1)

        return image, mask

    def _pad_to_size(
        self, tensor: torch.Tensor, target_size: Tuple[int, int]
    ) -> torch.Tensor:
        """Pad tensor to target size with zeros."""
        target_h, target_w = target_size

        if tensor.dim() == 3:  # Image [C, H, W]
            _, h, w = tensor.shape
        elif tensor.dim() == 2:  # Mask [H, W]
            h, w = tensor.shape
        else:
            raise ValueError(f"Unexpected tensor dimensions: {tensor.shape}")

        # Calculate padding
        pad_h = max(0, target_h - h)
        pad_w = max(0, target_w - w)

        # Pad equally on both sides
        pad_top = pad_h // 2
        pad_bottom = pad_h - pad_top
        pad_left = pad_w // 2
        pad_right = pad_w - pad_left

        # Apply padding (left, right, top, bottom)
        padded = F.pad(tensor, (pad_left, pad_right, pad_top, pad_bottom), value=0)

        # Crop if tensor is larger than target
        if tensor.dim() == 3:
            padded = padded[:, :target_h, :target_w]
        else:
            padded = padded[:target_h, :target_w]

        return padded

    def __len__(self) -> int:
        return len(self.image_paths)

    def __getitem__(self, idx: int) -> Tuple[torch.Tensor, torch.Tensor]:
        # Load image
        image_path = self.image_paths[idx]
        if self._is_geotiff(image_path):
            # Load GeoTIFF using rasterio
            with rasterio.open(image_path) as src:
                # Read as [C, H, W] format
                image = src.read().astype(np.float32)
                # Normalize image to [0, 1] range
                image = image / 255.0
        else:
            # Load standard image formats using PIL
            with Image.open(image_path) as img:
                # Convert to RGB if needed
                if img.mode != "RGB":
                    img = img.convert("RGB")
                # Convert to numpy array [H, W, C]
                image = np.array(img, dtype=np.float32)
                # Normalize to [0, 1] range
                image = image / 255.0
                # Convert to [C, H, W] format
                image = np.transpose(image, (2, 0, 1))

        # Handle different number of channels
        if image.shape[0] > self.num_channels:
            image = image[: self.num_channels]  # Keep only specified bands
        elif image.shape[0] < self.num_channels:
            # Pad with zeros if less than specified bands
            padded = np.zeros(
                (self.num_channels, image.shape[1], image.shape[2]),
                dtype=np.float32,
            )
            padded[: image.shape[0]] = image
            image = padded

        # Convert to CHW tensor
        image = torch.as_tensor(image, dtype=torch.float32)

        # Load label mask
        label_path = self.label_paths[idx]
        if self._is_geotiff(label_path):
            # Load GeoTIFF label using rasterio
            with rasterio.open(label_path) as src:
                label_mask = src.read(1).astype(np.int64)
        else:
            # Load standard image format label using PIL
            with Image.open(label_path) as img:
                # Convert to grayscale if needed
                if img.mode != "L":
                    img = img.convert("L")
                label_mask = np.array(img, dtype=np.int64)

        # Normalize mask values to expected class range [0, num_classes-1]
        # This handles cases where masks contain pixel values outside the expected range
        unique_vals = np.unique(label_mask)
        if self.num_classes == 2:
            # Binary segmentation: normalize any non-zero value to foreground (1)
            if unique_vals.max() > 1:
                label_mask = (label_mask > 0).astype(np.int64)
        else:
            # Multi-class segmentation:
            # Preserve class IDs and only clamp to valid range
            label_mask = np.clip(label_mask, 0, self.num_classes - 1)

        # Convert to tensor
        mask = torch.as_tensor(label_mask, dtype=torch.long)

        # Resize image and mask to target size if specified
        image, mask = self._resize_image_and_mask(image, mask)

        # Apply transforms if specified
        if self.transforms is not None:
            image, mask = self.transforms(image, mask)

        return image, mask


class SemanticTransforms:
    """Custom transforms for semantic segmentation."""

    def __init__(self, transforms: List[Callable]) -> None:
        self.transforms = transforms

    def __call__(
        self, image: torch.Tensor, mask: torch.Tensor
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        for t in self.transforms:
            image, mask = t(image, mask)
        return image, mask


class SemanticToTensor:
    """Convert numpy.ndarray to tensor for semantic segmentation."""

    def __call__(
        self, image: torch.Tensor, mask: torch.Tensor
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        return image, mask


class SemanticRandomHorizontalFlip:
    """Random horizontal flip transform for semantic segmentation."""

    def __init__(self, prob: float = 0.5) -> None:
        self.prob = prob

    def __call__(
        self, image: torch.Tensor, mask: torch.Tensor
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        if random.random() < self.prob:
            # Flip image and mask along width dimension
            image = torch.flip(image, dims=[2])
            mask = torch.flip(mask, dims=[1])
        return image, mask


class SemanticRandomVerticalFlip:
    """Random vertical flip transform for semantic segmentation."""

    def __init__(self, prob: float = 0.5) -> None:
        self.prob = prob

    def __call__(
        self, image: torch.Tensor, mask: torch.Tensor
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        if random.random() < self.prob:
            # Flip image and mask along height dimension
            image = torch.flip(image, dims=[1])
            mask = torch.flip(mask, dims=[0])
        return image, mask


class SemanticRandomRotation90:
    """Random 90-degree rotation transform for semantic segmentation."""

    def __init__(self, prob: float = 0.5) -> None:
        self.prob = prob

    def __call__(
        self, image: torch.Tensor, mask: torch.Tensor
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        if random.random() < self.prob:
            # Randomly rotate by 90, 180, or 270 degrees
            k = random.randint(1, 3)
            image = torch.rot90(image, k, dims=[1, 2])
            mask = torch.rot90(mask, k, dims=[0, 1])
        return image, mask


class SemanticBrightnessAdjustment:
    """Random brightness adjustment transform for semantic segmentation."""

    def __init__(
        self, brightness_range: Tuple[float, float] = (0.8, 1.2), prob: float = 0.5
    ) -> None:
        """
        Initialize brightness adjustment transform.

        Args:
            brightness_range: Tuple of (min, max) brightness factors.
            prob: Probability of applying the transform.
        """
        self.brightness_range = brightness_range
        self.prob = prob

    def __call__(
        self, image: torch.Tensor, mask: torch.Tensor
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        if random.random() < self.prob:
            # Apply random brightness adjustment
            factor = self.brightness_range[0] + random.random() * (
                self.brightness_range[1] - self.brightness_range[0]
            )
            image = torch.clamp(image * factor, 0, 1)
        return image, mask


class SemanticContrastAdjustment:
    """Random contrast adjustment transform for semantic segmentation."""

    def __init__(
        self, contrast_range: Tuple[float, float] = (0.8, 1.2), prob: float = 0.5
    ) -> None:
        """
        Initialize contrast adjustment transform.

        Args:
            contrast_range: Tuple of (min, max) contrast factors.
            prob: Probability of applying the transform.
        """
        self.contrast_range = contrast_range
        self.prob = prob

    def __call__(
        self, image: torch.Tensor, mask: torch.Tensor
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        if random.random() < self.prob:
            # Apply random contrast adjustment
            factor = self.contrast_range[0] + random.random() * (
                self.contrast_range[1] - self.contrast_range[0]
            )
            mean = image.mean(dim=(1, 2), keepdim=True)
            image = torch.clamp((image - mean) * factor + mean, 0, 1)
        return image, mask


def get_semantic_transform(train: bool) -> Any:
    """
    Get transforms for semantic segmentation data augmentation.

    This function returns default data augmentation transforms for training
    semantic segmentation models. The transforms include geometric transformations
    (horizontal/vertical flips, rotations) and photometric adjustments (brightness,
    contrast) that are commonly used in remote sensing tasks.

    Args:
        train (bool): Whether to include training-specific transforms.
            If True, applies augmentations (flips, rotations, brightness/contrast adjustments).
            If False, only converts to tensor (for validation).

    Returns:
        SemanticTransforms: Composed transforms.

    Example:
        >>> train_transform = get_semantic_transform(train=True)
        >>> val_transform = get_semantic_transform(train=False)
    """
    transforms = []
    transforms.append(SemanticToTensor())

    if train:
        # Geometric transforms - preserve spatial structure
        transforms.append(SemanticRandomHorizontalFlip(0.5))
        transforms.append(SemanticRandomVerticalFlip(0.5))
        transforms.append(SemanticRandomRotation90(0.5))

        # Photometric transforms - improve model robustness
        transforms.append(
            SemanticBrightnessAdjustment(brightness_range=(0.8, 1.2), prob=0.5)
        )
        transforms.append(
            SemanticContrastAdjustment(contrast_range=(0.8, 1.2), prob=0.5)
        )

    return SemanticTransforms(transforms)


def get_smp_model(
    architecture: str = "unet",
    encoder_name: str = "resnet34",
    encoder_weights: Optional[str] = "imagenet",
    in_channels: int = 3,
    classes: int = 2,
    activation: Optional[str] = None,
    **kwargs: Any,
) -> torch.nn.Module:
    """
    Get a segmentation model from segmentation-models-pytorch using the generic create_model function.

    Args:
        architecture (str): Model architecture (e.g., 'unet', 'deeplabv3', 'deeplabv3plus', 'fpn',
            'pspnet', 'linknet', 'manet', 'pan', 'upernet', etc.). Case insensitive.
        encoder_name (str): Encoder backbone name (e.g., 'resnet34', 'efficientnet-b0', 'mit_b0', etc.).
        encoder_weights (str): Encoder weights ('imagenet' or None).
        in_channels (int): Number of input channels.
        classes (int): Number of output classes.
        activation (str): Activation function for output layer.
        **kwargs: Additional arguments passed to smp.create_model().

    Returns:
        torch.nn.Module: Segmentation model.

    Note:
        This function uses smp.create_model() which supports all architectures available in
        segmentation-models-pytorch, making it future-proof for new model additions.
    """
    if not SMP_AVAILABLE:
        raise ImportError(
            "segmentation-models-pytorch is not installed. "
            "Please install it with: pip install segmentation-models-pytorch"
        )

    try:
        # Use the generic create_model function - supports all SMP architectures
        model = smp.create_model(
            arch=architecture,  # Case insensitive
            encoder_name=encoder_name,
            encoder_weights=encoder_weights,
            in_channels=in_channels,
            classes=classes,
            **kwargs,
        )

        # Apply activation if specified (note: activation is handled differently in create_model)
        if activation is not None:
            import warnings

            warnings.warn(
                "The 'activation' parameter is deprecated when using smp.create_model(). "
                "Apply activation manually after model creation if needed.",
                DeprecationWarning,
                stacklevel=2,
            )

        return model

    except Exception as e:
        # Provide helpful error message
        available_archs = []
        try:
            # Try to get available architectures from smp
            if hasattr(smp, "get_available_models"):
                available_archs = smp.get_available_models()
            else:
                available_archs = [
                    "unet",
                    "unetplusplus",
                    "manet",
                    "linknet",
                    "fpn",
                    "pspnet",
                    "deeplabv3",
                    "deeplabv3plus",
                    "pan",
                    "upernet",
                ]
        except Exception:
            available_archs = [
                "unet",
                "fpn",
                "deeplabv3plus",
                "pspnet",
                "linknet",
                "manet",
            ]

        raise ValueError(
            f"Failed to create model with architecture '{architecture}' and encoder '{encoder_name}'. "
            f"Error: {str(e)}. "
            f"Available architectures include: {', '.join(available_archs)}. "
            f"Please check the segmentation-models-pytorch documentation for supported combinations."
        )


def f1_score(
    pred: torch.Tensor,
    target: torch.Tensor,
    smooth: float = 1e-6,
    num_classes: Optional[int] = None,
) -> float:
    """
    Calculate F1 score (also known as Dice coefficient) for segmentation (binary or multi-class).

    Args:
        pred (torch.Tensor): Predicted mask (probabilities or logits) with shape [C, H, W] or [H, W].
        target (torch.Tensor): Ground truth mask with shape [H, W].
        smooth (float): Smoothing factor to avoid division by zero.
        num_classes (int, optional): Number of classes. If None, auto-detected.

    Returns:
        float: Mean F1 score across all classes.
    """
    # Convert predictions to class predictions
    if pred.dim() == 3:  # [C, H, W] format
        pred = torch.softmax(pred, dim=0)
        pred_classes = torch.argmax(pred, dim=0)
    elif pred.dim() == 2:  # [H, W] format
        pred_classes = pred
    else:
        raise ValueError(f"Unexpected prediction dimensions: {pred.shape}")

    # Auto-detect number of classes if not provided
    if num_classes is None:
        num_classes = max(pred_classes.max().item(), target.max().item()) + 1

    # Calculate F1 score for each class and average
    f1_scores = []
    for class_id in range(num_classes):
        pred_class = (pred_classes == class_id).float()
        target_class = (target == class_id).float()

        intersection = (pred_class * target_class).sum()
        union = pred_class.sum() + target_class.sum()

        if union > 0:
            f1 = (2.0 * intersection + smooth) / (union + smooth)
            f1_scores.append(f1.item())

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


def iou_coefficient(
    pred: torch.Tensor,
    target: torch.Tensor,
    smooth: float = 1e-6,
    num_classes: Optional[int] = None,
) -> float:
    """
    Calculate IoU coefficient for segmentation (binary or multi-class).

    Args:
        pred (torch.Tensor): Predicted mask (probabilities or logits) with shape [C, H, W] or [H, W].
        target (torch.Tensor): Ground truth mask with shape [H, W].
        smooth (float): Smoothing factor to avoid division by zero.
        num_classes (int, optional): Number of classes. If None, auto-detected.

    Returns:
        float: Mean IoU coefficient across all classes.
    """
    # Convert predictions to class predictions
    if pred.dim() == 3:  # [C, H, W] format
        pred = torch.softmax(pred, dim=0)
        pred_classes = torch.argmax(pred, dim=0)
    elif pred.dim() == 2:  # [H, W] format
        pred_classes = pred
    else:
        raise ValueError(f"Unexpected prediction dimensions: {pred.shape}")

    # Auto-detect number of classes if not provided
    if num_classes is None:
        num_classes = max(pred_classes.max().item(), target.max().item()) + 1

    # Calculate IoU for each class and average
    iou_scores = []
    for class_id in range(num_classes):
        pred_class = (pred_classes == class_id).float()
        target_class = (target == class_id).float()

        intersection = (pred_class * target_class).sum()
        union = pred_class.sum() + target_class.sum() - intersection

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

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


def precision_score(
    pred: torch.Tensor,
    target: torch.Tensor,
    smooth: float = 1e-6,
    num_classes: Optional[int] = None,
) -> float:
    """
    Calculate precision score for segmentation (binary or multi-class).

    Precision = TP / (TP + FP), where:
    - TP (True Positives): Correctly predicted positive pixels
    - FP (False Positives): Incorrectly predicted positive pixels

    Args:
        pred (torch.Tensor): Predicted mask (probabilities or logits) with shape [C, H, W] or [H, W].
        target (torch.Tensor): Ground truth mask with shape [H, W].
        smooth (float): Smoothing factor to avoid division by zero.
        num_classes (int, optional): Number of classes. If None, auto-detected.

    Returns:
        float: Mean precision score across all classes.
    """
    # Convert predictions to class predictions
    if pred.dim() == 3:  # [C, H, W] format
        pred = torch.softmax(pred, dim=0)
        pred_classes = torch.argmax(pred, dim=0)
    elif pred.dim() == 2:  # [H, W] format
        pred_classes = pred
    else:
        raise ValueError(f"Unexpected prediction dimensions: {pred.shape}")

    # Auto-detect number of classes if not provided
    if num_classes is None:
        num_classes = max(pred_classes.max().item(), target.max().item()) + 1

    # Calculate precision for each class and average
    precision_scores = []
    for class_id in range(num_classes):
        pred_class = (pred_classes == class_id).float()
        target_class = (target == class_id).float()

        true_positives = (pred_class * target_class).sum()
        predicted_positives = pred_class.sum()

        if predicted_positives > 0:
            precision = (true_positives + smooth) / (predicted_positives + smooth)
            precision_scores.append(precision.item())

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


def recall_score(
    pred: torch.Tensor,
    target: torch.Tensor,
    smooth: float = 1e-6,
    num_classes: Optional[int] = None,
) -> float:
    """
    Calculate recall score (also known as sensitivity) for segmentation (binary or multi-class).

    Recall = TP / (TP + FN), where:
    - TP (True Positives): Correctly predicted positive pixels
    - FN (False Negatives): Incorrectly predicted negative pixels

    Args:
        pred (torch.Tensor): Predicted mask (probabilities or logits) with shape [C, H, W] or [H, W].
        target (torch.Tensor): Ground truth mask with shape [H, W].
        smooth (float): Smoothing factor to avoid division by zero.
        num_classes (int, optional): Number of classes. If None, auto-detected.

    Returns:
        float: Mean recall score across all classes.
    """
    # Convert predictions to class predictions
    if pred.dim() == 3:  # [C, H, W] format
        pred = torch.softmax(pred, dim=0)
        pred_classes = torch.argmax(pred, dim=0)
    elif pred.dim() == 2:  # [H, W] format
        pred_classes = pred
    else:
        raise ValueError(f"Unexpected prediction dimensions: {pred.shape}")

    # Auto-detect number of classes if not provided
    if num_classes is None:
        num_classes = max(pred_classes.max().item(), target.max().item()) + 1

    # Calculate recall for each class and average
    recall_scores = []
    for class_id in range(num_classes):
        pred_class = (pred_classes == class_id).float()
        target_class = (target == class_id).float()

        true_positives = (pred_class * target_class).sum()
        actual_positives = target_class.sum()

        if actual_positives > 0:
            recall = (true_positives + smooth) / (actual_positives + smooth)
            recall_scores.append(recall.item())

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


def train_semantic_one_epoch(
    model: torch.nn.Module,
    optimizer: torch.optim.Optimizer,
    data_loader: DataLoader,
    device: torch.device,
    epoch: int,
    criterion: Any,
    print_freq: int = 10,
    verbose: bool = True,
) -> float:
    """
    Train the semantic segmentation model for one epoch.

    Args:
        model (torch.nn.Module): The model to train.
        optimizer (torch.optim.Optimizer): The optimizer to use.
        data_loader (torch.utils.data.DataLoader): DataLoader for training data.
        device (torch.device): Device to train on.
        epoch (int): Current epoch number.
        criterion: Loss function.
        print_freq (int): How often to print progress.
        verbose (bool): Whether to print detailed progress.

    Returns:
        float: Average loss for the epoch.
    """
    model.train()
    total_loss = 0
    num_batches = len(data_loader)

    start_time = time.time()

    for i, (images, targets) in enumerate(data_loader):
        # Move images and targets to device
        images = images.to(device)
        targets = targets.to(device)

        # Forward pass
        outputs = model(images)
        loss = criterion(outputs, targets)

        # Backward pass
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        # Track loss
        total_loss += loss.item()

        # Print progress
        if i % print_freq == 0:
            elapsed_time = time.time() - start_time
            if verbose:
                print(
                    f"Epoch: {epoch + 1}, Batch: {i + 1}/{num_batches}, Loss: {loss.item():.4f}, Time: {elapsed_time:.2f}s"
                )
            start_time = time.time()

    # Calculate average loss
    avg_loss = total_loss / num_batches
    return avg_loss


def evaluate_semantic(
    model: torch.nn.Module,
    data_loader: DataLoader,
    device: torch.device,
    criterion: Any,
    num_classes: int = 2,
) -> Dict[str, float]:
    """
    Evaluate the semantic segmentation model on the validation set.

    Args:
        model (torch.nn.Module): The model to evaluate.
        data_loader (torch.utils.data.DataLoader): DataLoader for validation data.
        device (torch.device): Device to evaluate on.
        criterion: Loss function.
        num_classes (int): Number of classes for evaluation metrics.

    Returns:
        dict: Evaluation metrics including loss, IoU, F1, precision, and recall.
    """
    model.eval()

    total_loss = 0
    f1_scores = []
    iou_scores = []
    precision_scores = []
    recall_scores = []
    num_batches = len(data_loader)

    with torch.no_grad():
        for images, targets in data_loader:
            # Move to device
            images = images.to(device)
            targets = targets.to(device)

            # Forward pass
            outputs = model(images)
            loss = criterion(outputs, targets)
            total_loss += loss.item()

            # Calculate metrics for each sample in the batch
            for pred, target in zip(outputs, targets):
                f1 = f1_score(pred, target, num_classes=num_classes)
                iou = iou_coefficient(pred, target, num_classes=num_classes)
                precision = precision_score(pred, target, num_classes=num_classes)
                recall = recall_score(pred, target, num_classes=num_classes)
                f1_scores.append(f1)
                iou_scores.append(iou)
                precision_scores.append(precision)
                recall_scores.append(recall)

    # Calculate metrics
    avg_loss = total_loss / num_batches
    avg_f1 = sum(f1_scores) / len(f1_scores) if f1_scores else 0
    avg_iou = sum(iou_scores) / len(iou_scores) if iou_scores else 0
    avg_precision = (
        sum(precision_scores) / len(precision_scores) if precision_scores else 0
    )
    avg_recall = sum(recall_scores) / len(recall_scores) if recall_scores else 0

    return {
        "loss": avg_loss,
        "F1": avg_f1,
        "IoU": avg_iou,
        "Precision": avg_precision,
        "Recall": avg_recall,
    }


def train_segmentation_model(
    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,
    early_stopping_patience: Optional[int] = None,
    train_transforms: Optional[Callable] = None,
    val_transforms: Optional[Callable] = None,
    **kwargs: Any,
) -> torch.nn.Module:
    """
    Train a semantic segmentation model for object detection using segmentation-models-pytorch.

    This function trains a semantic segmentation model for object detection (e.g., building detection)
    using models from the segmentation-models-pytorch library. Unlike instance segmentation (Mask R-CNN),
    this approach treats the task as pixel-level binary classification.

    Args:
        images_dir (str): Directory containing image GeoTIFF files (for 'directory' format),
            or root directory containing images/ subdirectory (for 'yolo' format),
            or directory containing images (for 'coco' format).
        labels_dir (str): Directory containing label GeoTIFF files (for 'directory' format),
            or path to COCO annotations JSON file (for 'coco' format),
            or not used (for 'yolo' format - labels are in images_dir/labels/).
        output_dir (str): Directory to save model checkpoints and results.
        input_format (str): Input data format - 'directory' (default), 'coco', or 'yolo'.
            - 'directory': Standard directory structure with separate images_dir and labels_dir
            - 'coco': COCO JSON format (labels_dir should be path to instances.json)
            - 'yolo': YOLO format (images_dir is root with images/ and labels/ subdirectories)
        architecture (str): Model architecture ('unet', 'deeplabv3', 'deeplabv3plus', 'fpn',
            'pspnet', 'linknet', 'manet'). Defaults to 'unet'.
        encoder_name (str): Encoder backbone name (e.g., 'resnet34', 'resnet50', 'efficientnet-b0').
            Defaults to 'resnet34'.
        encoder_weights (str): Encoder pretrained weights ('imagenet' or None). Defaults to 'imagenet'.
        num_channels (int): Number of input channels. Defaults to 3.
        num_classes (int): Number of output classes (typically 2 for binary segmentation). Defaults to 2.
        batch_size (int): Batch size for training. Defaults to 8.
        num_epochs (int): Number of training epochs. Defaults to 50.
        learning_rate (float): Initial learning rate. Defaults to 0.001.
        weight_decay (float): Weight decay for optimizer. Defaults to 1e-4.
        seed (int): Random seed for reproducibility. Defaults to 42.
        val_split (float): Fraction of data to use for validation (0-1). Defaults to 0.2.
        print_freq (int): Frequency of printing training progress. Defaults to 10.
        verbose (bool): If True, prints detailed training progress. Defaults to True.
        save_best_only (bool): If True, only saves the best model. Otherwise saves all checkpoints.
            Defaults to True.
        plot_curves (bool): If True, plots training curves. Defaults to False.
        device (torch.device): Device to train on. If None, uses CUDA if available.
        checkpoint_path (str, optional): Path to a checkpoint file to load for resuming training.
            If provided, will load model weights and optionally optimizer/scheduler state.
        resume_training (bool): If True and checkpoint_path is provided, will resume training
            from the checkpoint including optimizer and scheduler state. Defaults to False.
        target_size (tuple, optional): Target size (height, width) for standardizing images.
            If None, the function will automatically detect if images have varying sizes and set
            a default target_size of (512, 512) to prevent batching errors. To disable automatic
            resizing, set this parameter explicitly. Example: (512, 512). Defaults to None.
        resize_mode (str): How to handle size standardization when target_size is specified.
            'resize' - Resize images to target_size (may change aspect ratio)
            'pad' - Pad images to target_size (preserves aspect ratio). Defaults to 'resize'.
        num_workers (int): Number of workers for data loading. If None, uses 0 on macOS and Windows, 8 otherwise.
            Both image and mask should be torch.Tensor objects. The image tensor is expected to be in
            CHW format (channels, height, width), and the mask tensor in HW format (height, width).
            If None, uses default transforms (horizontal flip with 0.5 probability). Defaults to None.
        val_transforms (callable, optional): Custom transforms for validation data.
            Should be a callable that accepts (image, mask) tensors and returns transformed (image, mask).
            The image tensor is expected to be in CHW format (channels, height, width), and the mask tensor in HW format (height, width).
            Both image and mask should be torch.Tensor objects. If None, uses default transforms
            (horizontal flip with 0.5 probability). Defaults to None.
        val_transforms (callable, optional): Custom transforms for validation data.
            Should be a callable that accepts (image, mask) tensors and returns transformed (image, mask).
            If None, uses default transforms (no augmentation). Defaults to None.
        **kwargs: Additional arguments passed to smp.create_model().
    Returns:
        None: Model weights are saved to output_dir.

    Raises:
        ImportError: If segmentation-models-pytorch is not installed.
        FileNotFoundError: If input directories don't exist or contain no matching files.
    """
    import datetime

    if not SMP_AVAILABLE:
        raise ImportError(
            "segmentation-models-pytorch is not installed. "
            "Please install it with: pip install segmentation-models-pytorch"
        )

    # Set random seeds for reproducibility
    torch.manual_seed(seed)
    np.random.seed(seed)
    random.seed(seed)
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False

    # Create output directory
    os.makedirs(os.path.abspath(output_dir), exist_ok=True)

    # Get device
    if device is None:
        device = get_device()
    print(f"Using device: {device}")

    # Get all image and label files based on input format
    if input_format.lower() == "coco":
        # Parse COCO format annotations
        if verbose:
            print(f"Loading COCO format annotations from {labels_dir}")
        # For COCO format, labels_dir is path to instances.json
        # Labels are typically in a "labels" directory parallel to "annotations"
        coco_root = os.path.dirname(os.path.dirname(labels_dir))  # Go up two levels
        labels_directory = os.path.join(coco_root, "labels")
        image_files, label_files = parse_coco_annotations(
            labels_dir, images_dir, labels_directory
        )
    elif input_format.lower() == "yolo":
        # Parse YOLO format annotations
        if verbose:
            print(f"Loading YOLO format data from {images_dir}")
        image_files, label_files = parse_yolo_annotations(images_dir)
    else:
        # Default: directory format
        # Support multiple image formats: GeoTIFF, PNG, JPG, JPEG, TIF, TIFF
        image_extensions = (".tif", ".tiff", ".png", ".jpg", ".jpeg")
        label_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(label_extensions)
            ]
        )

        # Ensure matching files
        if len(image_files) != len(label_files):
            print("Warning: Number of image files and label files don't match!")
            # Find matching files by basename
            basenames = [os.path.basename(f) for f in image_files]
            label_files = [
                os.path.join(labels_dir, os.path.basename(f))
                for f in image_files
                if os.path.exists(os.path.join(labels_dir, os.path.basename(f)))
            ]
            image_files = [
                f
                for f, b in zip(image_files, basenames)
                if os.path.exists(os.path.join(labels_dir, b))
            ]
            print(f"Using {len(image_files)} matching files")

    print(f"Found {len(image_files)} image files and {len(label_files)} label files")

    if len(image_files) == 0:
        raise FileNotFoundError("No matching image and label files found")

    # Split data into train and validation sets
    train_imgs, val_imgs, train_labels, val_labels = train_test_split(
        image_files, label_files, test_size=val_split, random_state=seed
    )

    print(f"Training on {len(train_imgs)} images, validating on {len(val_imgs)} images")

    # Auto-detect image sizes and set target_size if needed
    if target_size is None:
        print("Checking image sizes for compatibility...")

        # Sample a few images to check size consistency
        sample_images = train_imgs[: min(5, len(train_imgs))]
        image_sizes = []

        for img_path in sample_images:
            try:
                if img_path.lower().endswith((".tif", ".tiff")):
                    with rasterio.open(img_path) as src:
                        height, width = src.height, src.width
                else:
                    with Image.open(img_path) as img:
                        width, height = img.size
                image_sizes.append((height, width))
            except Exception as e:
                print(f"Warning: Could not read image {img_path}: {e}")
                continue

        # Check if all images have the same size
        if len(image_sizes) == 0:
            print(
                "Warning: Could not read any sample images. Setting target_size to (512, 512) as a safe default."
            )
            target_size = (512, 512)
        else:
            unique_sizes = set(image_sizes)
            if len(unique_sizes) > 1:
                print(
                    f"Warning: Found images with different sizes: {list(unique_sizes)}"
                )
                print(
                    "Setting target_size to (512, 512) to standardize image dimensions."
                )
                print("This will resize all images to 512x512 pixels.")
                print("To use a different size, set target_size parameter explicitly.")
                target_size = (512, 512)
            else:
                print(f"All sampled images have the same size: {image_sizes[0]}")
                print("No resizing needed.")

    # Create datasets
    # Use custom transforms if provided, otherwise use default transforms
    train_transform = (
        train_transforms
        if train_transforms is not None
        else get_semantic_transform(train=True)
    )
    val_transform = (
        val_transforms
        if val_transforms is not None
        else get_semantic_transform(train=False)
    )

    train_dataset = SemanticSegmentationDataset(
        train_imgs,
        train_labels,
        transforms=train_transform,
        num_channels=num_channels,
        target_size=target_size,
        resize_mode=resize_mode,
        num_classes=num_classes,
    )
    val_dataset = SemanticSegmentationDataset(
        val_imgs,
        val_labels,
        transforms=val_transform,
        num_channels=num_channels,
        target_size=target_size,
        resize_mode=resize_mode,
        num_classes=num_classes,
    )

    # Create data loaders
    # Use num_workers=0 on macOS and Windows to avoid multiprocessing issues
    # Windows often has issues with multiprocessing in Jupyter notebooks
    # Increase num_workers for better data loading performance
    if num_workers is None:
        num_workers = 0 if platform.system() in ["Darwin", "Windows"] else 8

    try:
        train_loader = DataLoader(
            train_dataset,
            batch_size=batch_size,
            shuffle=True,
            num_workers=num_workers,
            pin_memory=True,
        )

        val_loader = DataLoader(
            val_dataset,
            batch_size=batch_size,
            shuffle=False,
            num_workers=num_workers,
            pin_memory=True,
        )

        # Test the data loader by loading one batch to catch size mismatch errors early
        print("Testing data loader...")
        try:
            next(iter(train_loader))
            print("Data loader test passed.")
        except RuntimeError as e:
            if "stack expects each tensor to be equal size" in str(e):
                raise RuntimeError(
                    "Images have different sizes and cannot be batched together. "
                    "Please set target_size parameter to standardize image dimensions. "
                    "Example: target_size=(512, 512). "
                    f"Original error: {str(e)}"
                ) from e
            else:
                raise

    except Exception as e:
        if "stack expects each tensor to be equal size" in str(e):
            raise RuntimeError(
                "Images have different sizes and cannot be batched together. "
                "Please set target_size parameter to standardize image dimensions. "
                "Example: target_size=(512, 512). "
                f"Original error: {str(e)}"
            ) from e
        else:
            raise

    # Initialize model
    model = get_smp_model(
        architecture=architecture,
        encoder_name=encoder_name,
        encoder_weights=encoder_weights,
        in_channels=num_channels,
        classes=num_classes,
        activation=None,  # We'll apply softmax later
        **kwargs,
    )
    model.to(device)

    # Enable multi-GPU training if multiple GPUs are available
    if torch.cuda.device_count() > 1:
        print(f"Using {torch.cuda.device_count()} GPUs for training")
        model = torch.nn.DataParallel(model)

    # Set up loss function (CrossEntropyLoss for multi-class, can also use F1Loss)
    criterion = torch.nn.CrossEntropyLoss()

    # Set up optimizer
    optimizer = torch.optim.Adam(
        model.parameters(), lr=learning_rate, weight_decay=weight_decay
    )

    # Set up learning rate scheduler
    lr_scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
        optimizer, mode="min", factor=0.5, patience=5
    )

    # Initialize tracking variables
    best_iou = 0
    train_losses = []
    val_losses = []
    val_ious = []
    val_f1s = []
    val_precisions = []
    val_recalls = []
    start_epoch = 0
    epochs_without_improvement = 0

    # Load checkpoint if provided
    if checkpoint_path is not None:
        if not os.path.exists(checkpoint_path):
            raise FileNotFoundError(f"Checkpoint file not found: {checkpoint_path}")

        print(f"Loading checkpoint from: {checkpoint_path}")
        try:
            checkpoint = torch.load(checkpoint_path, map_location=device)

            if isinstance(checkpoint, dict) and "model_state_dict" in checkpoint:
                # Load model state
                model.load_state_dict(checkpoint["model_state_dict"])

                if resume_training:
                    # Resume training from checkpoint
                    start_epoch = checkpoint.get("epoch", 0) + 1
                    best_iou = checkpoint.get("best_iou", 0)

                    # Load optimizer state if available
                    if "optimizer_state_dict" in checkpoint:
                        optimizer.load_state_dict(checkpoint["optimizer_state_dict"])

                    # Load scheduler state if available
                    if "scheduler_state_dict" in checkpoint:
                        lr_scheduler.load_state_dict(checkpoint["scheduler_state_dict"])

                    # Load training history if available
                    if "train_losses" in checkpoint:
                        train_losses = checkpoint["train_losses"]
                    if "val_losses" in checkpoint:
                        val_losses = checkpoint["val_losses"]
                    if "val_ious" in checkpoint:
                        val_ious = checkpoint["val_ious"]
                    if "val_f1s" in checkpoint:
                        val_f1s = checkpoint["val_f1s"]
                    # Also check for old val_dices format for backward compatibility
                    elif "val_dices" in checkpoint:
                        val_f1s = checkpoint["val_dices"]
                    if "val_precisions" in checkpoint:
                        val_precisions = checkpoint["val_precisions"]
                    if "val_recalls" in checkpoint:
                        val_recalls = checkpoint["val_recalls"]

                    print(f"Resuming training from epoch {start_epoch}")
                    print(f"Previous best IoU: {best_iou:.4f}")
                else:
                    print("Loaded model weights only (not resuming training state)")
            else:
                # Assume it's just model weights
                model.load_state_dict(checkpoint)
                print("Loaded model weights only")

        except Exception as e:
            raise RuntimeError(f"Failed to load checkpoint: {str(e)}")

    print(f"Starting training with {architecture} + {encoder_name}")
    print(f"Model parameters: {sum(p.numel() for p in model.parameters()):,}")
    if start_epoch > 0:
        print(f"Resuming from epoch {start_epoch}/{num_epochs}")

    # Training loop
    for epoch in range(start_epoch, num_epochs):
        # Train one epoch
        train_loss = train_semantic_one_epoch(
            model,
            optimizer,
            train_loader,
            device,
            epoch,
            criterion,
            print_freq,
            verbose,
        )
        train_losses.append(train_loss)

        # Evaluate on validation set
        eval_metrics = evaluate_semantic(
            model, val_loader, device, criterion, num_classes=num_classes
        )
        val_losses.append(eval_metrics["loss"])
        val_ious.append(eval_metrics["IoU"])
        val_f1s.append(eval_metrics["F1"])
        val_precisions.append(eval_metrics["Precision"])
        val_recalls.append(eval_metrics["Recall"])

        # Update learning rate
        lr_scheduler.step(eval_metrics["loss"])

        # Print metrics
        print(
            f"Epoch {epoch+1}/{num_epochs}: "
            f"Train Loss: {train_loss:.4f}, "
            f"Val Loss: {eval_metrics['loss']:.4f}, "
            f"Val IoU: {eval_metrics['IoU']:.4f}, "
            f"Val F1: {eval_metrics['F1']:.4f}, "
            f"Val Precision: {eval_metrics['Precision']:.4f}, "
            f"Val Recall: {eval_metrics['Recall']:.4f}"
        )

        # Save best model and check for early stopping
        if eval_metrics["IoU"] > best_iou:
            best_iou = eval_metrics["IoU"]
            epochs_without_improvement = 0
            print(f"Saving best model with IoU: {best_iou:.4f}")
            torch.save(model.state_dict(), os.path.join(output_dir, "best_model.pth"))
        else:
            epochs_without_improvement += 1
            if (
                early_stopping_patience is not None
                and epochs_without_improvement >= early_stopping_patience
            ):
                print(
                    f"\nEarly stopping triggered after {epochs_without_improvement} epochs without improvement"
                )
                print(f"Best validation IoU: {best_iou:.4f}")
                break

        # Save checkpoint every 10 epochs (if not save_best_only)
        if not save_best_only and ((epoch + 1) % 10 == 0 or epoch == num_epochs - 1):
            torch.save(
                {
                    "epoch": epoch,
                    "model_state_dict": model.state_dict(),
                    "optimizer_state_dict": optimizer.state_dict(),
                    "scheduler_state_dict": lr_scheduler.state_dict(),
                    "best_iou": best_iou,
                    "architecture": architecture,
                    "encoder_name": encoder_name,
                    "num_channels": num_channels,
                    "num_classes": num_classes,
                    "train_losses": train_losses,
                    "val_losses": val_losses,
                    "val_ious": val_ious,
                    "val_f1s": val_f1s,
                    "val_precisions": val_precisions,
                    "val_recalls": val_recalls,
                },
                os.path.join(output_dir, f"checkpoint_epoch_{epoch+1}.pth"),
            )

    # Save final model
    torch.save(model.state_dict(), os.path.join(output_dir, "final_model.pth"))

    # Save training history
    history = {
        "train_losses": train_losses,
        "val_losses": val_losses,
        "val_ious": val_ious,
        "val_f1s": val_f1s,
        "val_precisions": val_precisions,
        "val_recalls": val_recalls,
    }
    torch.save(history, os.path.join(output_dir, "training_history.pth"))

    # Save training summary
    with open(
        os.path.join(output_dir, "training_summary.txt"), "w", encoding="utf-8"
    ) as f:
        f.write(
            f"Training completed on: {datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n"
        )
        f.write(f"Architecture: {architecture}\n")
        f.write(f"Encoder: {encoder_name}\n")
        f.write(f"Total epochs: {num_epochs}\n")
        f.write(f"Best validation IoU: {best_iou:.4f}\n")
        f.write(f"Final validation IoU: {val_ious[-1]:.4f}\n")
        f.write(f"Final validation F1: {val_f1s[-1]:.4f}\n")
        f.write(f"Final validation Precision: {val_precisions[-1]:.4f}\n")
        f.write(f"Final validation Recall: {val_recalls[-1]:.4f}\n")
        f.write(f"Final validation loss: {val_losses[-1]:.4f}\n")

    print(f"Training complete! Best IoU: {best_iou:.4f}")
    print(f"Models saved to {output_dir}")

    # Plot training curves
    if plot_curves:
        try:
            plt.figure(figsize=(15, 5))

            plt.subplot(1, 3, 1)
            plt.plot(train_losses, label="Train Loss")
            plt.plot(val_losses, label="Val Loss")
            plt.title("Loss")
            plt.xlabel("Epoch")
            plt.ylabel("Loss")
            plt.legend()
            plt.grid(True)

            plt.subplot(1, 3, 2)
            plt.plot(val_ious, label="Val IoU")
            plt.title("IoU Score")
            plt.xlabel("Epoch")
            plt.ylabel("IoU")
            plt.legend()
            plt.grid(True)

            plt.subplot(1, 3, 3)
            plt.plot(val_f1s, label="Val F1")
            plt.title("F1 Score")
            plt.xlabel("Epoch")
            plt.ylabel("F1")
            plt.legend()
            plt.grid(True)

            plt.tight_layout()
            plt.savefig(
                os.path.join(output_dir, "training_curves.png"),
                dpi=150,
                bbox_inches="tight",
            )
            print(
                f"Training curves saved to {os.path.join(output_dir, 'training_curves.png')}"
            )
            plt.close()
        except Exception as e:
            print(f"Could not save training curves: {e}")


def semantic_inference_on_geotiff(
    model: torch.nn.Module,
    geotiff_path: str,
    output_path: str,
    window_size: int = 512,
    overlap: int = 256,
    batch_size: int = 4,
    num_channels: int = 3,
    num_classes: int = 2,
    device: Optional[torch.device] = None,
    probability_path: Optional[str] = None,
    probability_threshold: Optional[float] = None,
    save_class_probabilities: bool = False,
    quiet: bool = False,
    **kwargs: Any,
) -> Tuple[str, float]:
    """
    Perform semantic segmentation inference on a large GeoTIFF using a sliding window approach.

    Args:
        model (torch.nn.Module): Trained semantic segmentation model.
        geotiff_path (str): Path to input GeoTIFF file.
        output_path (str): Path to save output mask GeoTIFF.
        window_size (int): Size of sliding window for inference.
        overlap (int): Overlap between adjacent windows.
        batch_size (int): Batch size for inference.
        num_channels (int): Number of channels to use from the input image.
        num_classes (int): Number of classes in the model output.
        device (torch.device, optional): Device to run inference on.
        probability_path (str, optional): Path to save probability map. If provided,
            the normalized class probabilities will be saved as a multi-band raster.
        probability_threshold (float, optional): Probability threshold for binary classification.
            Only used when num_classes=2. If provided, pixels with class 1 probability >= threshold
            are classified as class 1, otherwise class 0. If None (default), uses argmax.
        save_class_probabilities (bool): If True and probability_path is provided, saves each
            class probability as a separate single-band file. Defaults to False.
        quiet (bool): If True, suppress progress bar. Defaults to False.
        **kwargs: Additional arguments.

    Returns:
        tuple: Tuple containing output path and inference time in seconds.
    """
    if device is None:
        device = get_device()

    # Put model in evaluation mode
    model.to(device)
    model.eval()

    # Open the GeoTIFF
    with rasterio.open(geotiff_path) as src:
        # Read metadata
        meta = src.meta
        height = src.height
        width = src.width

        # Update metadata for output raster
        out_meta = meta.copy()
        out_meta.update({"count": 1, "dtype": "uint8"})

        # Initialize accumulator arrays for multi-class probability blending
        # We'll accumulate probabilities for each class and then take argmax
        prob_accumulator = np.zeros((num_classes, height, width), dtype=np.float32)
        count_accumulator = np.zeros((height, width), dtype=np.float32)

        # Calculate steps
        steps_y = math.ceil((height - overlap) / (window_size - overlap))
        steps_x = math.ceil((width - overlap) / (window_size - overlap))
        last_y = height - window_size
        last_x = width - window_size

        total_windows = steps_y * steps_x
        if not quiet:
            print(f"Processing {total_windows} windows...")

        if not quiet:
            pbar = tqdm(total=total_windows)
        else:
            pbar = None

        batch_inputs = []
        batch_positions = []
        batch_count = 0

        start_time = time.time()

        for i in range(steps_y + 1):
            y = min(i * (window_size - overlap), last_y)
            y = max(0, y)

            if y > last_y and i > 0:
                continue

            for j in range(steps_x + 1):
                x = min(j * (window_size - overlap), last_x)
                x = max(0, x)

                if x > last_x and j > 0:
                    continue

                # Read window
                window = src.read(window=Window(x, y, window_size, window_size))

                if window.shape[1] == 0 or window.shape[2] == 0:
                    continue

                current_height = window.shape[1]
                current_width = window.shape[2]

                # Normalize and prepare input
                image = window.astype(np.float32) / 255.0

                # Handle different number of bands
                if image.shape[0] > num_channels:
                    image = image[:num_channels]
                elif image.shape[0] < num_channels:
                    padded = np.zeros(
                        (num_channels, current_height, current_width), dtype=np.float32
                    )
                    padded[: image.shape[0]] = image
                    image = padded

                # Convert to tensor
                image_tensor = torch.tensor(image, device=device)

                # Add to batch
                batch_inputs.append(image_tensor)
                batch_positions.append((y, x, current_height, current_width))
                batch_count += 1

                # Process batch
                if batch_count == batch_size or (i == steps_y and j == steps_x):
                    with torch.no_grad():
                        batch_tensor = torch.stack(batch_inputs)
                        outputs = model(batch_tensor)

                        # Apply softmax to get class probabilities
                        probs = torch.softmax(outputs, dim=1)

                    # Process each output in the batch
                    for idx, prob in enumerate(probs):
                        y_pos, x_pos, h, w = batch_positions[idx]

                        # Create weight matrix for blending
                        y_grid, x_grid = np.mgrid[0:h, 0:w]
                        dist_from_left = x_grid
                        dist_from_right = w - x_grid - 1
                        dist_from_top = y_grid
                        dist_from_bottom = h - y_grid - 1

                        edge_distance = np.minimum.reduce(
                            [
                                dist_from_left,
                                dist_from_right,
                                dist_from_top,
                                dist_from_bottom,
                            ]
                        )
                        edge_distance = np.minimum(edge_distance, overlap / 2)

                        # Avoid zero weights - use minimum weight of 0.1
                        weight = np.maximum(edge_distance / (overlap / 2), 0.1)

                        # For non-overlapping windows, use uniform weight
                        if overlap == 0:
                            weight = np.ones_like(weight)

                        # Convert probabilities to numpy [C, H, W]
                        prob_np = prob.cpu().numpy()

                        # Accumulate weighted probabilities for each class
                        y_slice = slice(y_pos, y_pos + h)
                        x_slice = slice(x_pos, x_pos + w)

                        # Add weighted probabilities for each class
                        for class_idx in range(num_classes):
                            prob_accumulator[class_idx, y_slice, x_slice] += (
                                prob_np[class_idx] * weight
                            )

                        # Update weight accumulator
                        count_accumulator[y_slice, x_slice] += weight

                    # Reset batch
                    batch_inputs = []
                    batch_positions = []
                    batch_count = 0
                    if pbar is not None:
                        pbar.update(len(probs))

        if pbar is not None:
            pbar.close()

        # Calculate final mask by taking argmax of accumulated probabilities
        mask = np.zeros((height, width), dtype=np.uint8)
        valid_pixels = count_accumulator > 0

        if np.any(valid_pixels):
            # Normalize accumulated probabilities by weights
            normalized_probs = np.zeros_like(prob_accumulator)
            for class_idx in range(num_classes):
                normalized_probs[class_idx, valid_pixels] = (
                    prob_accumulator[class_idx, valid_pixels]
                    / count_accumulator[valid_pixels]
                )

            # Apply threshold for binary classification or use argmax
            if probability_threshold is not None and num_classes == 2:
                # Use threshold: classify as class 1 if probability >= threshold
                mask[valid_pixels] = (
                    normalized_probs[1, valid_pixels] >= probability_threshold
                ).astype(np.uint8)
                if not quiet:
                    print(f"Using probability threshold: {probability_threshold}")
            else:
                # Take argmax to get final class predictions
                mask[valid_pixels] = np.argmax(
                    normalized_probs[:, valid_pixels], axis=0
                ).astype(np.uint8)

            # Check class distribution in predictions (summary only)
            unique_classes, class_counts = np.unique(
                mask[valid_pixels], return_counts=True
            )
            bg_ratio = np.sum(mask == 0) / mask.size
            if not quiet:
                print(
                    f"Predicted classes: {len(unique_classes)} classes, Background: {bg_ratio:.1%}"
                )

        inference_time = time.time() - start_time
        if not quiet:
            print(f"Inference completed in {inference_time:.2f} seconds")

        # Save output
        out_dir = os.path.abspath(os.path.dirname(output_path))
        os.makedirs(out_dir, exist_ok=True)
        with rasterio.open(output_path, "w", **out_meta) as dst:
            dst.write(mask, 1)

        if not quiet:
            print(f"Saved prediction to {output_path}")

        # Save probability map if requested
        if probability_path is not None:
            prob_dir = os.path.abspath(os.path.dirname(probability_path))
            os.makedirs(prob_dir, exist_ok=True)

            # Prepare probability output metadata
            prob_meta = meta.copy()
            prob_meta.update({"count": num_classes, "dtype": "float32"})

            # Save normalized probabilities as multi-band raster
            with rasterio.open(probability_path, "w", **prob_meta) as dst:
                for class_idx in range(num_classes):
                    # Normalize probabilities
                    prob_band = np.zeros((height, width), dtype=np.float32)
                    prob_band[valid_pixels] = (
                        prob_accumulator[class_idx, valid_pixels]
                        / count_accumulator[valid_pixels]
                    )
                    dst.write(prob_band, class_idx + 1)

            if not quiet:
                print(f"Saved probability map to {probability_path}")

            # Save individual class probabilities if requested
            if save_class_probabilities:
                # Prepare single-band metadata
                single_band_meta = meta.copy()
                single_band_meta.update({"count": 1, "dtype": "float32"})

                # Get base filename and extension
                prob_base = os.path.splitext(probability_path)[0]
                prob_ext = os.path.splitext(probability_path)[1]

                for class_idx in range(num_classes):
                    # Create filename for this class
                    class_prob_path = f"{prob_base}_class_{class_idx}{prob_ext}"

                    # Normalize probabilities
                    prob_band = np.zeros((height, width), dtype=np.float32)
                    prob_band[valid_pixels] = (
                        prob_accumulator[class_idx, valid_pixels]
                        / count_accumulator[valid_pixels]
                    )

                    # Save single-band file
                    with rasterio.open(class_prob_path, "w", **single_band_meta) as dst:
                        dst.write(prob_band, 1)

                    if not quiet:
                        print(
                            f"Saved class {class_idx} probability to {class_prob_path}"
                        )

        return output_path, inference_time


def semantic_inference_on_image(
    model: torch.nn.Module,
    image_path: str,
    output_path: str,
    window_size: int = 512,
    overlap: int = 256,
    batch_size: int = 4,
    num_channels: int = 3,
    num_classes: int = 2,
    device: Optional[torch.device] = None,
    binary_output: bool = True,
    probability_path: Optional[str] = None,
    probability_threshold: Optional[float] = None,
    save_class_probabilities: bool = False,
    quiet: bool = False,
    **kwargs: Any,
) -> Tuple[str, float]:
    """
    Perform semantic segmentation inference on a regular image (JPG, PNG, etc.) using a sliding window approach.

    Args:
        model (torch.nn.Module): Trained semantic segmentation model.
        image_path (str): Path to input image file (JPG, PNG, etc.).
        output_path (str): Path to save output mask image.
        window_size (int): Size of sliding window for inference.
        overlap (int): Overlap between adjacent windows.
        batch_size (int): Batch size for inference.
        num_channels (int): Number of channels to use from the input image.
        num_classes (int): Number of classes in the model output.
        device (torch.device, optional): Device to run inference on.
        binary_output (bool): If True, convert multi-class output to binary (class > 0).
        probability_path (str, optional): Path to save probability map. If provided,
            the normalized class probabilities will be saved as a multi-band raster.
        probability_threshold (float, optional): Probability threshold for binary classification.
            Only used when num_classes=2. If provided, pixels with class 1 probability >= threshold
            are classified as class 1, otherwise class 0. If None (default), uses argmax.
        save_class_probabilities (bool): If True and probability_path is provided, saves each
            class probability as a separate single-band file. Defaults to False.
        quiet (bool): If True, suppress progress bar. Defaults to False.
        **kwargs: Additional arguments.

    Returns:
        tuple: Tuple containing output path and inference time in seconds.
    """
    from PIL import Image

    if device is None:
        device = get_device()

    # Put model in evaluation mode
    model.to(device)
    model.eval()

    # Open the image using PIL
    with Image.open(image_path) as pil_img:
        # Convert to RGB if needed
        if pil_img.mode != "RGB":
            pil_img = pil_img.convert("RGB")

        # Convert to numpy array [H, W, C]
        img_array = np.array(pil_img, dtype=np.uint8)
        height, width = img_array.shape[:2]

        # Convert to [C, H, W] format like rasterio
        img_array = np.transpose(img_array, (2, 0, 1))

        if not quiet:
            print(f"Processing image: {width}x{height}")

        # Initialize accumulator arrays for multi-class probability blending
        prob_accumulator = np.zeros((num_classes, height, width), dtype=np.float32)
        count_accumulator = np.zeros((height, width), dtype=np.float32)

        # Calculate steps
        steps_y = math.ceil((height - overlap) / (window_size - overlap))
        steps_x = math.ceil((width - overlap) / (window_size - overlap))
        last_y = height - window_size
        last_x = width - window_size

        total_windows = steps_y * steps_x
        if not quiet:
            print(f"Processing {total_windows} windows...")

        if not quiet:
            pbar = tqdm(total=total_windows)
        else:
            pbar = None

        batch_inputs = []
        batch_positions = []
        batch_count = 0

        start_time = time.time()

        for i in range(steps_y + 1):
            y = min(i * (window_size - overlap), last_y)
            y = max(0, y)

            if y > last_y and i > 0:
                continue

            for j in range(steps_x + 1):
                x = min(j * (window_size - overlap), last_x)
                x = max(0, x)

                if x > last_x and j > 0:
                    continue

                # Extract window from image array
                y_end = min(y + window_size, height)
                x_end = min(x + window_size, width)
                window = img_array[:, y:y_end, x:x_end]

                if window.shape[1] == 0 or window.shape[2] == 0:
                    continue

                current_height = window.shape[1]
                current_width = window.shape[2]

                # Pad window to window_size if needed
                if current_height < window_size or current_width < window_size:
                    padded_window = np.zeros(
                        (window.shape[0], window_size, window_size), dtype=window.dtype
                    )
                    padded_window[:, :current_height, :current_width] = window
                    window = padded_window

                # Normalize and prepare input
                image = window.astype(np.float32) / 255.0

                # Handle different number of channels
                if image.shape[0] > num_channels:
                    image = image[:num_channels]
                elif image.shape[0] < num_channels:
                    padded = np.zeros(
                        (num_channels, image.shape[1], image.shape[2]), dtype=np.float32
                    )
                    padded[: image.shape[0]] = image
                    image = padded

                # Convert to tensor
                image_tensor = torch.tensor(image, device=device)

                # Add to batch
                batch_inputs.append(image_tensor)
                batch_positions.append((y, x, current_height, current_width))
                batch_count += 1

                # Process batch
                if batch_count == batch_size or (i == steps_y and j == steps_x):
                    with torch.no_grad():
                        batch_tensor = torch.stack(batch_inputs)
                        outputs = model(batch_tensor)

                        # Apply softmax to get class probabilities
                        probs = torch.softmax(outputs, dim=1)

                    # Process each output in the batch
                    for idx, prob in enumerate(probs):
                        y_pos, x_pos, h, w = batch_positions[idx]

                        # Create weight matrix for blending
                        y_grid, x_grid = np.mgrid[0:h, 0:w]
                        dist_from_left = x_grid
                        dist_from_right = w - x_grid - 1
                        dist_from_top = y_grid
                        dist_from_bottom = h - y_grid - 1

                        edge_distance = np.minimum.reduce(
                            [
                                dist_from_left,
                                dist_from_right,
                                dist_from_top,
                                dist_from_bottom,
                            ]
                        )
                        edge_distance = np.minimum(edge_distance, overlap / 2)

                        # Avoid zero weights - use minimum weight of 0.1
                        weight = np.maximum(edge_distance / (overlap / 2), 0.1)

                        # For non-overlapping windows, use uniform weight
                        if overlap == 0:
                            weight = np.ones_like(weight)

                        # Convert probabilities to numpy [C, H, W] - crop to actual size
                        prob_np = prob.cpu().numpy()[:, :h, :w]

                        # Accumulate weighted probabilities for each class
                        y_slice = slice(y_pos, y_pos + h)
                        x_slice = slice(x_pos, x_pos + w)

                        # Add weighted probabilities for each class
                        for class_idx in range(num_classes):
                            prob_accumulator[class_idx, y_slice, x_slice] += (
                                prob_np[class_idx] * weight
                            )

                        # Update weight accumulator
                        count_accumulator[y_slice, x_slice] += weight

                    # Reset batch
                    batch_inputs = []
                    batch_positions = []
                    batch_count = 0
                    if pbar is not None:
                        pbar.update(len(probs))

        if pbar is not None:
            pbar.close()

        # Calculate final mask by taking argmax of accumulated probabilities
        mask = np.zeros((height, width), dtype=np.uint8)
        valid_pixels = count_accumulator > 0

        if np.any(valid_pixels):
            # Normalize accumulated probabilities by weights
            normalized_probs = np.zeros_like(prob_accumulator)
            for class_idx in range(num_classes):
                normalized_probs[class_idx, valid_pixels] = (
                    prob_accumulator[class_idx, valid_pixels]
                    / count_accumulator[valid_pixels]
                )

            # Apply threshold for binary classification or use argmax
            if probability_threshold is not None and num_classes == 2:
                # Use threshold: classify as class 1 if probability >= threshold
                mask[valid_pixels] = (
                    normalized_probs[1, valid_pixels] >= probability_threshold
                ).astype(np.uint8)
                if not quiet:
                    print(f"Using probability threshold: {probability_threshold}")
            else:
                # Take argmax to get final class predictions
                mask[valid_pixels] = np.argmax(
                    normalized_probs[:, valid_pixels], axis=0
                ).astype(np.uint8)

            # Check class distribution in predictions before binary conversion
            unique_classes, class_counts = np.unique(mask, return_counts=True)
            # Convert numpy types to regular Python types for cleaner output
            class_distribution = {
                int(cls): int(count) for cls, count in zip(unique_classes, class_counts)
            }
            if not quiet:
                print(f"Raw predicted classes and counts: {class_distribution}")

            # Convert to binary if requested and num_classes == 2
            if binary_output and num_classes == 2:
                # For binary segmentation, convert class 1 to 255 (white) and class 0 to 0 (black)
                # Use proper thresholding to ensure only 0 and 255 values
                binary_mask = np.zeros_like(mask)
                binary_mask[mask > 0] = 255
                mask = binary_mask

                # Final check
                unique_classes, class_counts = np.unique(mask, return_counts=True)
                # Convert numpy types to regular Python types for cleaner output
                binary_distribution = {
                    int(cls): int(count)
                    for cls, count in zip(unique_classes, class_counts)
                }
                if not quiet:
                    print(f"Binary predicted classes and counts: {binary_distribution}")

        inference_time = time.time() - start_time
        if not quiet:
            print(f"Inference completed in {inference_time:.2f} seconds")

        # Save output as image
        # For binary masks, use PNG to avoid JPEG compression artifacts
        if binary_output and num_classes == 2:
            # Change extension to PNG if binary output to preserve exact values
            output_path_png = os.path.splitext(output_path)[0] + ".png"
            output_img = Image.fromarray(mask, mode="L")
            out_dir = os.path.abspath(os.path.dirname(output_path))
            os.makedirs(out_dir, exist_ok=True)
            output_img.save(output_path_png)
            if not quiet:
                print(
                    f"Saved binary prediction to {output_path_png} (PNG format to preserve exact values)"
                )

            # Also save the original requested format for compatibility
            if output_path != output_path_png:
                output_img.save(output_path)
                print(f"Also saved to {output_path} (may have compression artifacts)")
        else:
            output_img = Image.fromarray(mask, mode="L")
            output_img.save(output_path)
            if not quiet:
                print(f"Saved prediction to {output_path}")

        # Save probability map if requested
        if probability_path is not None:
            prob_dir = os.path.abspath(os.path.dirname(probability_path))
            os.makedirs(prob_dir, exist_ok=True)

            # For regular images, we'll save as a multi-channel TIFF
            # since we need to preserve floating point values
            import rasterio
            from rasterio.transform import from_bounds

            # Create a simple affine transform (identity transform for pixel coordinates)
            transform = from_bounds(0, 0, width, height, width, height)

            # Prepare probability output metadata
            prob_meta = {
                "driver": "GTiff",
                "height": height,
                "width": width,
                "count": num_classes,
                "dtype": "float32",
                "transform": transform,
            }

            # Save normalized probabilities as multi-band raster
            with rasterio.open(probability_path, "w", **prob_meta) as dst:
                for class_idx in range(num_classes):
                    # Normalize probabilities
                    prob_band = np.zeros((height, width), dtype=np.float32)
                    prob_band[valid_pixels] = normalized_probs[class_idx, valid_pixels]
                    dst.write(prob_band, class_idx + 1)

            if not quiet:
                print(f"Saved probability map to {probability_path}")

            # Save individual class probabilities if requested
            if save_class_probabilities:
                # Prepare single-band metadata
                single_band_meta = {
                    "driver": "GTiff",
                    "height": height,
                    "width": width,
                    "count": 1,
                    "dtype": "float32",
                    "transform": transform,
                }

                # Get base filename and extension
                prob_base = os.path.splitext(probability_path)[0]
                prob_ext = os.path.splitext(probability_path)[1]

                for class_idx in range(num_classes):
                    # Create filename for this class
                    class_prob_path = f"{prob_base}_class_{class_idx}{prob_ext}"

                    # Normalize probabilities
                    prob_band = np.zeros((height, width), dtype=np.float32)
                    prob_band[valid_pixels] = normalized_probs[class_idx, valid_pixels]

                    # Save single-band file
                    with rasterio.open(class_prob_path, "w", **single_band_meta) as dst:
                        dst.write(prob_band, 1)

                    if not quiet:
                        print(
                            f"Saved class {class_idx} probability to {class_prob_path}"
                        )

        return output_path, inference_time


def semantic_segmentation(
    input_path: str,
    output_path: str,
    model_path: str,
    architecture: str = "unet",
    encoder_name: str = "resnet34",
    num_channels: int = 3,
    num_classes: int = 2,
    window_size: int = 512,
    overlap: int = 256,
    batch_size: int = 4,
    device: Optional[torch.device] = None,
    probability_path: Optional[str] = None,
    probability_threshold: Optional[float] = None,
    save_class_probabilities: bool = False,
    quiet: bool = False,
    **kwargs: Any,
) -> None:
    """
    Perform semantic segmentation on an image file using a trained model.

    This function automatically detects the input format and uses the appropriate
    inference method for either GeoTIFF files or regular image formats (JPG, PNG, etc.).

    Args:
        input_path (str): Path to input image file (GeoTIFF, JPG, PNG, etc.).
        output_path (str): Path to save output mask file.
        model_path (str): Path to trained model weights.
        architecture (str): Model architecture used for training.
        encoder_name (str): Encoder backbone name used for training.
        num_channels (int): Number of channels in the input image and model.
        num_classes (int): Number of classes in the model.
        window_size (int): Size of sliding window for inference.
        overlap (int): Overlap between adjacent windows.
        batch_size (int): Batch size for inference.
        device (torch.device, optional): Device to run inference on.
        probability_path (str, optional): Path to save probability map. If provided,
            the normalized class probabilities will be saved as a multi-band raster
            where each band contains probabilities for each class.
        probability_threshold (float, optional): Probability threshold for binary classification.
            Only used when num_classes=2. If provided, pixels with class 1 probability >= threshold
            are classified as class 1, otherwise class 0. If None (default), uses argmax.
            Must be between 0 and 1.
        save_class_probabilities (bool): If True and probability_path is provided, saves each
            class probability as a separate single-band file. Files will be named like
            "probability_class_0.tif", "probability_class_1.tif", etc. in the same directory
            as probability_path. Defaults to False.
        quiet (bool): If True, suppress progress bar. Defaults to False.
        **kwargs: Additional arguments.

    Returns:
        None: Output mask is saved to output_path.
    """
    if device is None:
        device = get_device()

    # Detect file format based on extension
    input_ext = os.path.splitext(input_path)[1].lower()
    is_geotiff = input_ext in [".tif", ".tiff", ".jp2", ".img"]
    formats = {
        ".tif": "GeoTIFF",
        ".tiff": "GeoTIFF",
        ".jp2": "JP2OpenJPEG",
        ".img": "IMG",
    }

    if not quiet:
        print(
            f"Input file format: {formats[input_ext] if is_geotiff else 'Regular image'} ({input_ext})"
        )

    # Load model
    model = get_smp_model(
        architecture=architecture,
        encoder_name=encoder_name,
        encoder_weights=None,  # We're loading trained weights
        in_channels=num_channels,
        classes=num_classes,
        activation=None,
    )

    if not os.path.exists(model_path):
        try:
            model_path = download_model_from_hf(model_path)
        except Exception as e:
            raise FileNotFoundError(f"Model file not found: {model_path}")

    # Load state dict and handle DataParallel module prefix
    state_dict = torch.load(model_path, map_location=device)

    # Remove 'module.' prefix if present (from DataParallel training)
    if any(key.startswith("module.") for key in state_dict.keys()):
        state_dict = {
            key.replace("module.", ""): value for key, value in state_dict.items()
        }

    model.load_state_dict(state_dict)
    model.to(device)
    model.eval()

    # Validate probability_threshold
    if probability_threshold is not None:
        if not (0 <= probability_threshold <= 1):
            raise ValueError("probability_threshold must be between 0 and 1")
        if num_classes != 2:
            raise ValueError(
                "probability_threshold is only supported for binary classification (num_classes=2)"
            )

    # Use appropriate inference function based on file format
    if is_geotiff:
        semantic_inference_on_geotiff(
            model=model,
            geotiff_path=input_path,
            output_path=output_path,
            window_size=window_size,
            overlap=overlap,
            batch_size=batch_size,
            num_channels=num_channels,
            num_classes=num_classes,
            device=device,
            probability_path=probability_path,
            probability_threshold=probability_threshold,
            save_class_probabilities=save_class_probabilities,
            quiet=quiet,
            **kwargs,
        )
    else:
        # Create output directory if it doesn't exist
        os.makedirs(os.path.abspath(os.path.dirname(output_path)), exist_ok=True)

        semantic_inference_on_image(
            model=model,
            image_path=input_path,
            output_path=output_path,
            window_size=window_size,
            overlap=overlap,
            batch_size=batch_size,
            num_channels=num_channels,
            num_classes=num_classes,
            device=device,
            binary_output=True,  # Convert to binary output for better visualization
            probability_path=probability_path,
            probability_threshold=probability_threshold,
            save_class_probabilities=save_class_probabilities,
            quiet=quiet,
            **kwargs,
        )


def semantic_segmentation_batch(
    input_dir: str,
    output_dir: str,
    model_path: str,
    architecture: str = "unet",
    encoder_name: str = "resnet34",
    num_channels: int = 3,
    num_classes: int = 2,
    window_size: int = 512,
    overlap: int = 256,
    batch_size: int = 4,
    device: Optional[torch.device] = None,
    filenames: Optional[List[str]] = None,
    quiet: bool = False,
    **kwargs: Any,
) -> None:
    """
    Perform semantic segmentation on a batch of images from an input directory.

    This function processes all images in a directory and saves the results to an output directory.
    It automatically detects the input format and uses the appropriate inference method for either
    GeoTIFF files or regular image formats (JPG, PNG, etc.). For GeoTIFF inputs, outputs are saved
    as GeoTIFF. For other formats, outputs are saved as PNG to preserve exact values.

    Args:
        input_dir (str): Directory containing input image files to process.
        output_dir (str): Directory to save output mask files.
        model_path (str): Path to trained model weights.
        architecture (str): Model architecture used for training. Defaults to "unet".
        encoder_name (str): Encoder backbone name used for training. Defaults to "resnet34".
        num_channels (int): Number of channels in the input image and model. Defaults to 3.
        num_classes (int): Number of classes in the model. Defaults to 2.
        window_size (int): Size of sliding window for inference. Defaults to 512.
        overlap (int): Overlap between adjacent windows. Defaults to 256.
        batch_size (int): Batch size for inference. Defaults to 4.
        device (torch.device, optional): Device to run inference on. If None, uses CUDA if available.
        filenames (list, optional): List of output filenames. If None, defaults to
            "<input_filename>_mask.<ext>" for each input file where <ext> is "tif" for GeoTIFF
            inputs and "png" for other formats. If provided, must match the number of input files.
        quiet (bool): If True, suppress progress bar. Defaults to False.
        **kwargs: Additional arguments passed to the inference functions.

    Returns:
        None: Output masks are saved to output_dir.

    Raises:
        FileNotFoundError: If input_dir doesn't exist or contains no supported image files.
        ValueError: If filenames is provided but doesn't match the number of input files.
    """
    if device is None:
        device = get_device()

    # Check if input directory exists
    if not os.path.exists(input_dir):
        raise FileNotFoundError(f"Input directory does not exist: {input_dir}")

    # Create output directory if it doesn't exist
    os.makedirs(os.path.abspath(output_dir), exist_ok=True)

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

    if len(image_files) == 0:
        raise FileNotFoundError(f"No supported image files found in {input_dir}")

    print(f"Found {len(image_files)} image files to process")

    # Load model once for all images
    model = get_smp_model(
        architecture=architecture,
        encoder_name=encoder_name,
        encoder_weights=None,  # We're loading trained weights
        in_channels=num_channels,
        classes=num_classes,
        activation=None,
    )

    if not os.path.exists(model_path):
        try:
            model_path = download_model_from_hf(model_path)
        except Exception as e:
            raise FileNotFoundError(f"Model file not found: {model_path}")

    # Load state dict and handle DataParallel module prefix
    state_dict = torch.load(model_path, map_location=device)

    # Remove 'module.' prefix if present (from DataParallel training)
    if any(key.startswith("module.") for key in state_dict.keys()):
        state_dict = {
            key.replace("module.", ""): value for key, value in state_dict.items()
        }

    model.load_state_dict(state_dict)
    model.to(device)
    model.eval()

    # Generate output filenames if not provided
    if filenames is None:
        filenames = []
        for image_file in image_files:
            base_name = os.path.splitext(os.path.basename(image_file))[0]
            input_ext = os.path.splitext(image_file)[1].lower()

            # Use GeoTIFF extension for GeoTIFF inputs, PNG for others
            if input_ext in [".tif", ".tiff"]:
                output_ext = ".tif"
            else:
                output_ext = ".png"

            output_filename = f"{base_name}_mask{output_ext}"
            filenames.append(os.path.join(output_dir, output_filename))
    else:
        # Validate filenames list
        if len(filenames) != len(image_files):
            raise ValueError(
                f"Number of filenames ({len(filenames)}) must match number of input files ({len(image_files)})"
            )

    # Process each image
    for i, (input_path, output_path) in enumerate(zip(image_files, filenames)):
        print(
            f"Processing file {i + 1}/{len(image_files)}: {os.path.basename(input_path)}"
        )

        # Detect file format based on extension
        input_ext = os.path.splitext(input_path)[1].lower()
        is_geotiff = input_ext in [".tif", ".tiff"]

        try:
            # Use appropriate inference function based on file format
            if is_geotiff:
                semantic_inference_on_geotiff(
                    model=model,
                    geotiff_path=input_path,
                    output_path=output_path,
                    window_size=window_size,
                    overlap=overlap,
                    batch_size=batch_size,
                    num_channels=num_channels,
                    num_classes=num_classes,
                    device=device,
                    quiet=quiet,
                    **kwargs,
                )
            else:
                semantic_inference_on_image(
                    model=model,
                    image_path=input_path,
                    output_path=output_path,
                    window_size=window_size,
                    overlap=overlap,
                    batch_size=batch_size,
                    num_channels=num_channels,
                    num_classes=num_classes,
                    device=device,
                    binary_output=True,  # Convert to binary output for better visualization
                    quiet=quiet,
                    **kwargs,
                )
        except Exception as e:
            print(f"Error processing {input_path}: {str(e)}")
            continue

    print(f"Batch processing completed. Results saved to {output_dir}")


def train_instance_segmentation_model(
    images_dir: str,
    labels_dir: str,
    output_dir: str,
    input_format: str = "directory",
    num_classes: int = 2,
    num_channels: int = 3,
    batch_size: int = 4,
    num_epochs: int = 10,
    learning_rate: float = 0.005,
    seed: int = 42,
    val_split: float = 0.2,
    visualize: bool = False,
    device: Optional[torch.device] = None,
    verbose: bool = True,
    instance_labels: bool = False,
    **kwargs: Any,
) -> torch.nn.Module:
    """
    Train an instance segmentation model using Mask R-CNN.

    This is a wrapper function for train_MaskRCNN_model with clearer naming.

    Args:
        images_dir (str): Directory containing image GeoTIFF files (for 'directory' format),
            or root directory containing images/ subdirectory (for 'yolo' format),
            or directory containing images (for 'coco' format).
        labels_dir (str): Directory containing label GeoTIFF files (for 'directory' format),
            or path to COCO annotations JSON file (for 'coco' format),
            or not used (for 'yolo' format - labels are in images_dir/labels/).
        output_dir (str): Directory to save model checkpoints and results.
        input_format (str): Input data format - 'directory' (default), 'coco', or 'yolo'.
            - 'directory': Standard directory structure with separate images_dir and labels_dir
            - 'coco': COCO JSON format (labels_dir should be path to instances.json)
            - 'yolo': YOLO format (images_dir is root with images/ and labels/ subdirectories)
        num_classes (int): Number of classes (including background). Defaults to 2.
        num_channels (int): Number of input channels. Defaults to 3.
        batch_size (int): Batch size for training. Defaults to 4.
        num_epochs (int): Number of training epochs. Defaults to 10.
        learning_rate (float): Initial learning rate. Defaults to 0.005.
        seed (int): Random seed for reproducibility. Defaults to 42.
        val_split (float): Fraction of data to use for validation (0-1). Defaults to 0.2.
        visualize (bool): Whether to generate visualizations. Defaults to False.
        device (torch.device): Device to train on. If None, uses CUDA if available.
        verbose (bool): If True, prints detailed training progress. Defaults to True.
        instance_labels (bool): If True, treat label mask pixel values as
            pre-assigned instance IDs instead of running connected-component
            labeling. Use this when label masks already encode unique integers
            per object (e.g., Fields of The World dataset). Defaults to False.
        **kwargs: Additional arguments passed to train_MaskRCNN_model.

    Returns:
        None: Model weights are saved to output_dir.
    """
    # Create model with the specified number of classes
    model = get_instance_segmentation_model(
        num_classes=num_classes, num_channels=num_channels, pretrained=True
    )

    return train_MaskRCNN_model(
        images_dir=images_dir,
        labels_dir=labels_dir,
        output_dir=output_dir,
        input_format=input_format,
        num_channels=num_channels,
        model=model,
        batch_size=batch_size,
        num_epochs=num_epochs,
        learning_rate=learning_rate,
        seed=seed,
        val_split=val_split,
        visualize=visualize,
        device=device,
        verbose=verbose,
        instance_labels=instance_labels,
        **kwargs,
    )


def instance_segmentation(
    input_path: str,
    output_path: str,
    model_path: str,
    window_size: int = 512,
    overlap: int = 256,
    confidence_threshold: float = 0.5,
    nms_threshold: float = 0.3,
    batch_size: int = 4,
    num_channels: int = 3,
    num_classes: int = 2,
    class_names: Optional[List[str]] = None,
    vectorize: bool = False,
    vector_path: Optional[str] = None,
    use_mask_geometry: bool = True,
    simplify_tolerance: float = 0.0,
    device: Optional[torch.device] = None,
    **kwargs: Any,
) -> Dict[str, Any]:
    """
    Perform instance segmentation on a GeoTIFF using a pre-trained Mask
    R-CNN model.

    Saves three raster files (instance IDs, class labels, confidence scores)
    and optionally a vector file with polygon geometries and attributes.

    Args:
        input_path (str): Path to input GeoTIFF file.
        output_path (str): Path to save the instance ID raster. Class and
            score rasters are saved alongside with ``_class`` and ``_score``
            suffixes.
        model_path (str): Path to trained model weights.
        window_size (int): Size of sliding window for inference. Defaults
            to 512.
        overlap (int): Overlap between adjacent windows. Defaults to 256.
        confidence_threshold (float): Confidence threshold for predictions
            (0-1). Defaults to 0.5.
        nms_threshold (float): IoU threshold for non-maximum suppression.
            Defaults to 0.3.
        batch_size (int): Batch size for inference. Defaults to 4.
        num_channels (int): Number of channels in the input image and model.
            Defaults to 3.
        num_classes (int): Number of classes (including background).
            Defaults to 2.
        class_names (list, optional): List of class names where index 0 is
            the background class. Used for the vector output class_name
            attribute.
        vectorize (bool): If True, convert detections to a vector file with
            polygon geometries and attributes (geometry, class_id,
            class_name, score, instance_id, area_pixels). Defaults to False.
        vector_path (str, optional): Path to save the vector output. If None
            and vectorize is True, derives the path from output_path with a
            ``.gpkg`` extension.
        use_mask_geometry (bool): If True, use the instance mask polygons as
            geometry. If False, use bounding box rectangles. Only used when
            vectorize is True. Defaults to True.
        simplify_tolerance (float): Tolerance for polygon simplification in
            georeferenced units. Only used when vectorize is True and
            use_mask_geometry is True. Defaults to 0.0.
        device (torch.device): Device to run inference on. If None, uses
            CUDA if available.
        **kwargs: Additional arguments passed to
            instance_segmentation_inference_on_geotiff.

    Returns:
        dict: Dictionary with keys:
            - instance (str): Path to the instance ID raster.
            - class_label (str): Path to the class label raster.
            - score (str): Path to the confidence score raster.
            - vector (str or None): Path to the vector file, if vectorize
              is True.
            - inference_time (float): Inference time in seconds.
            - detections (list): List of detection dicts with keys: mask
              (window-sized np.ndarray), mask_offset (tuple of y, x, h, w),
              score, box, label, instance_id.
    """
    from .object_detect import detections_to_geodataframe

    # Create model with the specified number of classes
    model = get_instance_segmentation_model(
        num_classes=num_classes, num_channels=num_channels, pretrained=True
    )

    # Load the trained model
    if device is None:
        device = get_device()

    # Load state dict and handle DataParallel module prefix
    state_dict = torch.load(model_path, map_location=device)

    # Remove 'module.' prefix if present (from DataParallel training)
    if any(key.startswith("module.") for key in state_dict.keys()):
        state_dict = {
            key.replace("module.", ""): value for key, value in state_dict.items()
        }

    model.load_state_dict(state_dict)
    model.to(device)

    # Run inference
    output_paths, inference_time, detections = (
        instance_segmentation_inference_on_geotiff(
            model=model,
            geotiff_path=input_path,
            output_path=output_path,
            window_size=window_size,
            overlap=overlap,
            confidence_threshold=confidence_threshold,
            nms_threshold=nms_threshold,
            batch_size=batch_size,
            num_channels=num_channels,
            device=device,
            **kwargs,
        )
    )

    result = {
        "instance": output_paths["instance"],
        "class_label": output_paths["class_label"],
        "score": output_paths["score"],
        "vector": None,
        "inference_time": inference_time,
        "detections": detections,
    }

    # Optionally create vector output
    if vectorize and len(detections) > 0:
        if vector_path is None:
            base, _ = os.path.splitext(output_path)
            vector_path = f"{base}.gpkg"

        gdf = detections_to_geodataframe(
            detections=detections,
            geotiff_path=input_path,
            class_names=class_names,
            use_mask_geometry=use_mask_geometry,
            simplify_tolerance=simplify_tolerance,
        )
        gdf.to_file(vector_path, driver="GPKG")
        result["vector"] = vector_path
        print(f"Saved vector output to {vector_path}")

    return result


def instance_segmentation_batch(
    input_dir: str,
    output_dir: str,
    model_path: str,
    window_size: int = 512,
    overlap: int = 256,
    confidence_threshold: float = 0.5,
    nms_threshold: float = 0.3,
    batch_size: int = 4,
    num_channels: int = 3,
    num_classes: int = 2,
    device: Optional[torch.device] = None,
    **kwargs: Any,
) -> None:
    """
    Perform instance segmentation on multiple GeoTIFF files using a pre-trained
    Mask R-CNN model.

    Args:
        input_dir (str): Directory containing input GeoTIFF files.
        output_dir (str): Directory to save output mask GeoTIFF files.
        model_path (str): Path to trained model weights.
        window_size (int): Size of sliding window for inference. Defaults to 512.
        overlap (int): Overlap between adjacent windows. Defaults to 256.
        confidence_threshold (float): Confidence threshold for predictions
            (0-1). Defaults to 0.5.
        nms_threshold (float): IoU threshold for non-maximum suppression.
            Defaults to 0.3.
        batch_size (int): Batch size for inference. Defaults to 4.
        num_channels (int): Number of channels in the input image and model.
            Defaults to 3.
        num_classes (int): Number of classes (including background).
            Defaults to 2.
        device (torch.device): Device to run inference on. If None, uses
            CUDA if available.
        **kwargs: Additional arguments passed to
            instance_segmentation_inference_on_geotiff.

    Returns:
        None: Output masks are saved to output_dir.
    """
    # Create model with the specified number of classes
    model = get_instance_segmentation_model(
        num_classes=num_classes, num_channels=num_channels, pretrained=True
    )

    # Load the trained model
    if device is None:
        device = get_device()

    # Load state dict and handle DataParallel module prefix
    state_dict = torch.load(model_path, map_location=device)

    # Remove 'module.' prefix if present (from DataParallel training)
    if any(key.startswith("module.") for key in state_dict.keys()):
        state_dict = {
            key.replace("module.", ""): value for key, value in state_dict.items()
        }

    model.load_state_dict(state_dict)
    model.to(device)

    # Process all GeoTIFF files in the input directory
    import glob

    input_files = glob.glob(os.path.join(input_dir, "*.tif")) + glob.glob(
        os.path.join(input_dir, "*.tiff")
    )

    if not input_files:
        print(f"No GeoTIFF files found in {input_dir}")
        return

    # Create output directory if it doesn't exist
    os.makedirs(output_dir, exist_ok=True)

    print(f"Processing {len(input_files)} files...")

    for input_file in input_files:
        try:
            # Generate output filename
            base_name = os.path.splitext(os.path.basename(input_file))[0]
            output_file = os.path.join(output_dir, f"{base_name}_instances.tif")

            print(f"Processing {input_file}...")

            # Run instance segmentation inference
            instance_segmentation_inference_on_geotiff(
                model=model,
                geotiff_path=input_file,
                output_path=output_file,
                window_size=window_size,
                overlap=overlap,
                confidence_threshold=confidence_threshold,
                nms_threshold=nms_threshold,
                batch_size=batch_size,
                num_channels=num_channels,
                device=device,
                **kwargs,
            )

            print(f"Saved result to {output_file}")

        except Exception as e:
            print(f"Error processing {input_file}: {str(e)}")
            continue

    print(f"Batch processing completed. Results saved to {output_dir}")


def lightly_train_model(
    data_dir: str,
    output_dir: str,
    model: str = "torchvision/resnet50",
    method: str = "dinov2_distillation",
    epochs: int = 100,
    batch_size: int = 64,
    learning_rate: float = 1e-4,
    **kwargs: Any,
) -> str:
    """
    Train a model using Lightly Train for self-supervised pretraining.

    Args:
        data_dir (str): Directory containing unlabeled images for training.
        output_dir (str): Directory to save training outputs and model checkpoints.
        model (str): Model architecture to train. Supports models from torchvision,
            timm, ultralytics, etc. Default is "torchvision/resnet50".
        method (str): Self-supervised learning method. Options include:
            - "simclr": Works with CNN models (ResNet, EfficientNet, etc.)
            - "dino": Works with both CNNs and ViTs
            - "dinov2": Requires ViT models only
            - "dinov2_distillation": Requires ViT models only (recommended for ViTs)
            Default is "dinov2_distillation".
        epochs (int): Number of training epochs. Default is 100.
        batch_size (int): Batch size for training. Default is 64.
        learning_rate (float): Learning rate for training. Default is 1e-4.
        **kwargs: Additional arguments passed to lightly_train.train().

    Returns:
        str: Path to the exported model file.

    Raises:
        ImportError: If lightly-train is not installed.
        ValueError: If data_dir does not exist, is empty, or incompatible model/method.

    Note:
        Model/Method compatibility:
        - CNN models (ResNet, EfficientNet): Use "simclr" or "dino"
        - ViT models: Use "dinov2", "dinov2_distillation", or "dino"

    Example:
        >>> # For CNN models (ResNet, EfficientNet)
        >>> model_path = lightly_train_model(
        ...     data_dir="path/to/unlabeled/images",
        ...     output_dir="path/to/output",
        ...     model="torchvision/resnet50",
        ...     method="simclr",  # Use simclr for CNNs
        ...     epochs=50
        ... )
        >>> # For ViT models
        >>> model_path = lightly_train_model(
        ...     data_dir="path/to/unlabeled/images",
        ...     output_dir="path/to/output",
        ...     model="timm/vit_base_patch16_224",
        ...     method="dinov2",  # dinov2 requires ViT
        ...     epochs=50
        ... )
    """
    if not LIGHTLY_TRAIN_AVAILABLE:
        raise ImportError(
            "lightly-train is not installed. Please install it with: "
            "pip install lightly-train"
        )

    if not os.path.exists(data_dir):
        raise ValueError(f"Data directory does not exist: {data_dir}")

    # Check if data directory contains images
    image_extensions = ["*.jpg", "*.jpeg", "*.png", "*.tif", "*.tiff", "*.bmp"]
    image_files = []
    for ext in image_extensions:
        image_files.extend(glob.glob(os.path.join(data_dir, "**", ext), recursive=True))

    if not image_files:
        raise ValueError(f"No image files found in {data_dir}")

    # Validate model/method compatibility
    is_vit_model = "vit" in model.lower() or "vision_transformer" in model.lower()

    if method in ["dinov2", "dinov2_distillation"] and not is_vit_model:
        raise ValueError(
            f"Method '{method}' requires a Vision Transformer (ViT) model, but got '{model}'.\n"
            f"Solutions:\n"
            f"  1. Use a ViT model: model='timm/vit_base_patch16_224'\n"
            f"  2. Use a CNN-compatible method: method='simclr' or method='dino'\n"
            f"\nFor CNN models (ResNet, EfficientNet), use 'simclr' or 'dino'.\n"
            f"For ViT models, use 'dinov2', 'dinov2_distillation', or 'dino'."
        )

    print(f"Found {len(image_files)} images in {data_dir}")
    print(f"Starting self-supervised pretraining with {method} method...")
    print(f"Model: {model}")

    # Create output directory
    os.makedirs(output_dir, exist_ok=True)

    # Detect if running in notebook environment and set appropriate configuration
    def is_notebook():
        try:
            from IPython import get_ipython

            if get_ipython() is not None:
                return True
        except (ImportError, NameError):
            pass
        return False

    # Force single-device training in notebooks to avoid DDP strategy issues
    if is_notebook():
        # Only override if not explicitly set by user
        if "accelerator" not in kwargs:
            # Use CPU in notebooks to avoid DDP incompatibility
            # Users can still override by passing accelerator='gpu'
            kwargs["accelerator"] = "cpu"
        if "devices" not in kwargs:
            kwargs["devices"] = 1  # Force single device

    # Train the model using Lightly Train
    lightly_train.train(
        out=output_dir,
        data=data_dir,
        model=model,
        method=method,
        epochs=epochs,
        batch_size=batch_size,
        **kwargs,
    )

    # Return path to the exported model
    exported_model_path = os.path.join(
        output_dir, "exported_models", "exported_last.pt"
    )

    if os.path.exists(exported_model_path):
        print(
            f"Model training completed. Exported model saved to: {exported_model_path}"
        )
        return exported_model_path
    else:
        # Check for alternative export paths
        possible_paths = [
            os.path.join(output_dir, "exported_models", "exported_best.pt"),
            os.path.join(output_dir, "checkpoints", "last.ckpt"),
        ]

        for path in possible_paths:
            if os.path.exists(path):
                print(f"Model training completed. Exported model saved to: {path}")
                return path

        print(f"Model training completed. Output saved to: {output_dir}")
        return output_dir


def load_lightly_pretrained_model(
    model_path: str,
    model_architecture: str = "torchvision/resnet50",
    device: str = None,
) -> torch.nn.Module:
    """
    Load a pretrained model from Lightly Train.

    Args:
        model_path (str): Path to the pretrained model file (.pt format).
        model_architecture (str): Architecture of the model to load.
            Default is "torchvision/resnet50".
        device (str): Device to load the model on. If None, uses CPU.

    Returns:
        torch.nn.Module: Loaded pretrained model ready for fine-tuning.

    Raises:
        FileNotFoundError: If model_path does not exist.
        ImportError: If required libraries are not available.

    Example:
        >>> model = load_lightly_pretrained_model(
        ...     model_path="path/to/pretrained_model.pt",
        ...     model_architecture="torchvision/resnet50",
        ...     device="cuda"
        ... )
        >>> # Fine-tune the model with your existing training pipeline
    """
    if not os.path.exists(model_path):
        raise FileNotFoundError(f"Model file not found: {model_path}")

    print(f"Loading pretrained model from: {model_path}")

    # Load the model based on architecture
    if model_architecture.startswith("torchvision/"):
        model_name = model_architecture.replace("torchvision/", "")

        # Import the model from torchvision
        if hasattr(torchvision.models, model_name):
            model = getattr(torchvision.models, model_name)()
        else:
            raise ValueError(f"Unknown torchvision model: {model_name}")

    elif model_architecture.startswith("timm/"):
        try:
            import timm

            model_name = model_architecture.replace("timm/", "")
            model = timm.create_model(model_name)
        except ImportError:
            raise ImportError(
                "timm is required for TIMM models. Install with: pip install timm"
            )

    else:
        # For other architectures, try to import from torchvision as default
        try:
            model = getattr(torchvision.models, model_architecture)()
        except AttributeError:
            raise ValueError(f"Unsupported model architecture: {model_architecture}")

    # Load the pretrained weights
    try:
        state_dict = torch.load(model_path, map_location=device, weights_only=True)
    except TypeError:
        # For backward compatibility with older PyTorch versions
        state_dict = torch.load(model_path, map_location=device)
    model.load_state_dict(state_dict)

    print(f"Successfully loaded pretrained model: {model_architecture}")
    return model


def lightly_embed_images(
    data_dir: str,
    model_path: str,
    output_path: str,
    model_architecture: str = None,  # Deprecated, kept for backwards compatibility
    batch_size: int = 64,
    **kwargs: Any,
) -> str:
    """
    Generate embeddings for images using a Lightly Train pretrained model.

    Args:
        data_dir (str): Directory containing images to embed.
        model_path (str): Path to the pretrained model checkpoint file (.ckpt).
        output_path (str): Path to save the embeddings (as .pt file).
        model_architecture (str): Architecture of the pretrained model (deprecated,
            kept for backwards compatibility but not used). The model architecture
            is automatically loaded from the checkpoint.
        batch_size (int): Batch size for embedding generation. Default is 64.
        **kwargs: Additional arguments passed to lightly_train.embed().
            Supported kwargs include: image_size, num_workers, accelerator, etc.

    Returns:
        str: Path to the saved embeddings file.

    Raises:
        ImportError: If lightly-train is not installed.
        FileNotFoundError: If data_dir or model_path does not exist.

    Note:
        The model_path should point to a .ckpt file from the training output,
        typically located at: output_dir/checkpoints/last.ckpt

    Example:
        >>> embeddings_path = lightly_embed_images(
        ...     data_dir="path/to/images",
        ...     model_path="output_dir/checkpoints/last.ckpt",
        ...     output_path="embeddings.pt",
        ...     batch_size=32
        ... )
        >>> print(f"Embeddings saved to: {embeddings_path}")
    """
    if not LIGHTLY_TRAIN_AVAILABLE:
        raise ImportError(
            "lightly-train is not installed. Please install it with: "
            "pip install lightly-train"
        )

    if not os.path.exists(data_dir):
        raise FileNotFoundError(f"Data directory does not exist: {data_dir}")

    if not os.path.exists(model_path):
        raise FileNotFoundError(f"Model file does not exist: {model_path}")

    print(f"Generating embeddings for images in: {data_dir}")
    print(f"Using pretrained model: {model_path}")

    output_dir = os.path.dirname(output_path)
    if output_dir:
        os.makedirs(output_dir, exist_ok=True)

    # Generate embeddings using Lightly Train
    # Note: model_architecture is not used - it's inferred from the checkpoint
    lightly_train.embed(
        out=output_path,
        data=data_dir,
        checkpoint=model_path,
        batch_size=batch_size,
        **kwargs,
    )

    print(f"Embeddings saved to: {output_path}")
    return output_path
