Implement initial version of SegFormer training pipeline with dataset parsing and model training functionalities. Added Dockerfile for environment setup, utility scripts for parsing and training, and Gradio interface for user interaction.

Dockerfile

@@ -0,0 +1,37 @@
+FROM pytorch/pytorch:2.8.0-cuda12.9-cudnn9-runtime
+
+ENV PYTHONUNBUFFERED=1
+
+RUN useradd -m -u 1000 user
+WORKDIR /app
+
+# Install system dependencies for building Python packages
+RUN apt-get update && apt-get install -y \
+    build-essential \
+    libffi-dev \
+    libssl-dev \
+    ffmpeg \
+    libsm6 \
+    libxext6 \
+    libmagic1 \
+    && rm -rf /var/lib/apt/lists/*
+
+
+# Create virtual environment for Sly (keep isolated)
+RUN python -m venv /app/.venv-sly
+RUN /app/.venv-sly/bin/pip install --upgrade pip
+COPY --chown=user requirements-sly.txt requirements-sly.txt
+RUN /app/.venv-sly/bin/pip install -r requirements-sly.txt
+
+# Install Gradio and other dependencies
+RUN pip install --upgrade pip
+COPY --chown=user requirements.txt requirements.txt
+RUN pip install -r requirements.txt
+
+
+# Copy the rest of the app
+COPY --chown=user . .
+
+ENV GRADIO_SERVER_NAME="0.0.0.0"
+
+CMD ["python", "app.py"]

app.py

@@ -0,0 +1,346 @@
+import os
+import shutil
+import gradio as gr
+import pandas as pd
+
+from utils import run_supervisely_parser, train_model
+
+
+def run_pipeline(
+    base_model_zip,
+    supervisely_project_zip,
+    train_ratio,
+    seed,
+    data_percent,
+    batch_size,
+    num_epochs,
+    learning_rate,
+    image_width,
+    image_height,
+    early_stopping,
+    validate_every,
+    pr=gr.Progress(track_tqdm=True),
+):
+
+    # Parsing
+    yield (
+        gr.update(interactive=False),  # run button
+        gr.update(  # status textbox
+            value="Parsing Supervisely project ...",
+            visible=True,
+        ),
+        gr.update(visible=False),  # model download button
+        gr.update(visible=False),  # output tab
+        gr.update(value=None),  # train IoU plot
+        gr.update(value=None),  # val IoU plot
+        gr.update(value=None),  # metrics table
+    )
+
+    dataset_dir = run_supervisely_parser(
+        project_path=supervisely_project_zip,
+        train_ratio=train_ratio,
+        seed=seed,
+    )
+
+    # Training
+    yield (
+        gr.update(interactive=False),  # run button
+        gr.update(  # status textbox
+            value="Starting model training...",
+        ),
+        gr.update(visible=False),  # model download button
+        gr.update(visible=False),  # output tab
+        gr.update(value=None),  # train IoU plot
+        gr.update(value=None),  # val IoU plot
+        gr.update(value=None),  # metrics table
+    )
+
+    best_model, metrics, dice = train_model(
+        dataset_dir,
+        base_model_zip,
+        image_width,
+        image_height,
+        batch_size,
+        data_percent,
+        num_epochs,
+        learning_rate,
+        early_stopping,
+        validate_every,
+    )
+
+    # Saving model
+    yield (
+        gr.update(interactive=False),  # run button
+        gr.update(  # status textbox
+            value="Saving best model...",
+        ),
+        gr.update(visible=False),  # model download button
+        gr.update(visible=False),  # output tab
+        gr.update(value=None),  # train IoU plot
+        gr.update(value=None),  # val IoU plot
+        gr.update(value=None),  # metrics table
+    )
+    best_model_dir = os.path.join(
+        os.path.dirname(base_model_zip),
+        "best_model",
+    )
+    best_model.save_pretrained(best_model_dir)
+
+    best_model_zip_path = shutil.make_archive(
+        base_name=best_model_dir,
+        format="zip",
+        root_dir=best_model_dir,
+    )
+
+    metrics_df = pd.DataFrame(metrics)
+
+    initial_epoch_metrics = metrics_df.iloc[0]
+    final_epoch_metrics = metrics_df.iloc[-1]
+
+    # metrics comparison table use epoch 0 as before and final as after
+    metrics_comparison_df = pd.DataFrame(
+        {
+            "Metric": ["Accuracy", "IoU", "Loss", "Dice"],
+            "Before": [
+                initial_epoch_metrics["val_acc"],
+                initial_epoch_metrics["val_iou"],
+                initial_epoch_metrics["val_loss"],
+                dice[0],
+            ],
+            "After": [
+                final_epoch_metrics["val_acc"],
+                final_epoch_metrics["val_iou"],
+                final_epoch_metrics["val_loss"],
+                dice[1],
+            ],
+        }
+    )
+
+    yield (
+        gr.update(interactive=True),  # run button
+        gr.update(visible=False),  # status textbox
+        gr.update(  # model download button
+            value=best_model_zip_path,
+            visible=True,
+        ),
+        gr.update(visible=True),  # output tab
+        gr.update(value=metrics_df),  # train IoU plot
+        gr.update(value=metrics_df),  # val IoU plot
+        gr.update(
+            value=metrics_comparison_df,
+            visible=True,
+        ),
+    )
+
+
+def _toggle_run_btn(base_model, project):
+    """Enable run button only when both required files are selected."""
+    ready = bool(base_model and project)
+    return gr.update(interactive=ready)
+
+
+with gr.Blocks(title="SegFormer Training & Dataset Pipeline") as demo:
+    gr.Markdown(
+        "# SegFormer Training Pipeline\n"
+        "Upload your base model and Supervisely project, "
+        "tweak parsing & training hyperparameters, then click "
+        "**Run Training**."
+    )
+
+    with gr.Row():
+        base_model_zip = gr.File(
+            label="Base PyTorch Model (.zip)",
+            file_types=[".zip"],
+            file_count="single",
+        )
+        supervisely_project_zip = gr.File(
+            label="Supervisely Project (.zip)",
+            file_types=[".zip"],
+            file_count="single",
+        )
+
+    with gr.Tab("Training"):
+        gr.Markdown("Adjust training hyperparameters.")
+        with gr.Row():
+            data_percent = gr.Slider(
+                minimum=1,
+                maximum=100,
+                step=1,
+                value=100,
+                label="Data Percent (%) used for training",
+            )
+            batch_size = gr.Number(
+                value=32,
+                label="Batch Size (samples/step)",
+                precision=0,
+                minimum=1,
+            )
+            num_epochs = gr.Number(
+                value=60,
+                label="Epochs (max passes)",
+                precision=0,
+                minimum=1,
+            )
+        with gr.Row():
+            learning_rate = gr.Number(
+                value=5e-5,
+                label="Learning Rate",
+                minimum=0.0,
+                maximum=1.0,
+            )
+            image_width = gr.Number(
+                value=640,
+                label="Image Width (px)",
+                precision=0,
+                minimum=1,
+            )
+            image_height = gr.Number(
+                value=640,
+                label="Image Height (px)",
+                precision=0,
+                minimum=1,
+            )
+        with gr.Row():
+            early_stopping = gr.Number(
+                value=3,
+                label="Early Stopping Patience (epochs w/o improvement)",
+                precision=0,
+                minimum=0,
+            )
+            validate_every = gr.Number(
+                value=1,
+                label="Validate Every (epochs)",
+                precision=0,
+                minimum=0,
+            )
+
+    with gr.Tab("Dataset Parsing"):
+        gr.Markdown("Configure how the dataset is split and seeded.")
+        with gr.Row():
+            train_ratio = gr.Slider(
+                minimum=0.1,
+                maximum=0.95,
+                step=0.01,
+                value=0.8,
+                label="Train Split Ratio (rest used for validation)",
+            )
+            seed = gr.Number(
+                value=42,
+                label="Random Seed (reproducibility)",
+                precision=0,
+            )
+
+    with gr.Accordion("Parameter Help", open=False):
+        gr.Markdown(
+            """
+            **Base PyTorch Model (.zip)**: Archive containing a folder with
+            weights and configuration file.\n
+            **Supervisely Project (.zip)**: Archive containing Exported
+            Supervisely project
+            containing images and annotation JSONs.\n
+            **Train Split Ratio**: Fraction of dataset used for training;
+            remainder becomes validation.\n
+            **Random Seed**: Controls shuffling for reproducible splits &
+            training.\n
+            **Data Percent**: Subsample percentage of training split (use
+            <100 for quick experiments).\n
+            **Batch Size**: Samples processed before each optimizer step.\n
+            **Epochs**: Maximum complete passes over the (subsampled)
+            training set.\n
+            **Learning Rate**: Initial optimizer step size.\n
+            **Image Width / Height**: Target spatial size for preprocessing
+            (resize/crop).\n
+            **Early Stopping Patience**: Stop after this many validation
+            checks without improvement.\n
+            **Validate Every**: Run validation after this many epochs.\n
+            """
+        )
+
+    run_btn = gr.Button(
+        "Run Training",
+        variant="primary",
+        interactive=False,
+    )
+    status = gr.Textbox(
+        show_label=False,
+        visible=False,
+    )
+
+    with gr.Tab("Results", visible=False) as output_tab:
+
+        model_download_btn = gr.DownloadButton(
+            label="Download Trained Model (.zip)",
+            value=None,
+            visible=False,
+        )
+
+        # table to show before and after accuracy and iou
+        metrics_table = gr.DataFrame(
+            label="Metrics Comparison",
+            interactive=False,
+            wrap=True,
+        )
+
+        with gr.Row():
+            train_iou_plot = gr.LinePlot(
+                label="Training IoU",
+                x="epoch",
+                y="train_iou",
+                x_title="Epoch",
+                y_title="IoU",
+                height=400,
+            )
+            val_iou_plot = gr.LinePlot(
+                label="Validation IoU",
+                x="epoch",
+                y="val_iou",
+                x_title="Epoch",
+                y_title="IoU",
+                height=400,
+            )
+
+    # Enable run button only when both archives provided
+    base_model_zip.change(
+        _toggle_run_btn,
+        inputs=[base_model_zip, supervisely_project_zip],
+        outputs=run_btn,
+    )
+    supervisely_project_zip.change(
+        _toggle_run_btn,
+        inputs=[base_model_zip, supervisely_project_zip],
+        outputs=run_btn,
+    )
+
+    # Click handler
+    run_btn.click(
+        run_pipeline,
+        inputs=[
+            base_model_zip,
+            supervisely_project_zip,
+            train_ratio,
+            seed,
+            data_percent,
+            batch_size,
+            num_epochs,
+            learning_rate,
+            image_width,
+            image_height,
+            early_stopping,
+            validate_every,
+        ],
+        outputs=[
+            run_btn,
+            status,
+            model_download_btn,
+            output_tab,
+            train_iou_plot,
+            val_iou_plot,
+            metrics_table,
+        ],
+        show_progress_on=status,
+        scroll_to_output=True,
+    )
+
+
+if __name__ == "__main__":
+    demo.launch()

requirements-sly.txt

@@ -0,0 +1,4 @@
+supervisely>=6.73.418
+numpy>=1.26.4
+pillow>=10.2.0
+tqdm>=4.67.1

requirements.txt

@@ -0,0 +1,6 @@
+torch==2.8.0
+gradio>=5.42.0
+pillow>=11.3.0
+transformers>=4.55.2
+tqdm>=4.67.1
+evaluate>=0.4.5

scripts/supervisely_parser.py

@@ -0,0 +1,221 @@
+"""
+Supervisely Parser Script
+
+This script parses Supervisely projects and converts them to a format
+suitable for training segmentation models. It extracts class information,
+creates train/validation splits, and converts annotations to indexed
+color masks.
+"""
+
+import os
+import json
+import random
+import shutil
+import argparse
+import numpy as np
+from PIL import Image
+from tqdm import tqdm
+
+try:
+    import supervisely as sly
+    from supervisely import Annotation
+except ImportError as e:
+    print(f"Failed to import supervisely: {e}")
+    print(
+        "Please ensure that the 'supervisely' package is installed and "
+        "compatible with your environment."
+    )
+    raise
+
+
+def extract_class_info(project, output_dir):
+    """Extract class information from the project metadata."""
+    id2label = {}
+    id2color = {}
+
+    for obj in project.meta.obj_classes:
+        id_str, _, label = obj.name.partition(". ")
+        if not label or not id_str.isdigit():
+            continue
+        index = int(id_str) - 1
+
+        id2label[index] = label
+        id2color[index] = obj.color
+
+    # Save class mappings
+    with open(f"{output_dir}/id2label.json", "w") as f:
+        json.dump(id2label, f, sort_keys=True, indent=2)
+    with open(f"{output_dir}/id2color.json", "w") as f:
+        json.dump(id2color, f, sort_keys=True, indent=2)
+
+    label2id = {v: k for k, v in id2label.items()}
+    return id2label, id2color, label2id
+
+
+def create_output_directories(output_dir):
+    """Create necessary output directories."""
+    os.makedirs(f"{output_dir}/images/training", exist_ok=True)
+    os.makedirs(f"{output_dir}/annotations/training", exist_ok=True)
+    os.makedirs(f"{output_dir}/images/validation", exist_ok=True)
+    os.makedirs(f"{output_dir}/annotations/validation", exist_ok=True)
+
+
+def calculate_split_counts(datasets, train_ratio=0.8):
+    """Calculate the number of items for training and validation."""
+    total_items = 0
+    for dataset in datasets:
+        total_items += len(dataset.get_items_names())
+
+    train_items = int(total_items * train_ratio)
+    val_items = total_items - train_items
+
+    print(
+        f"Total items: {total_items}\n"
+        f"Train items: {train_items}\n"
+        f"Validation items: {val_items}"
+    )
+
+    return train_items, val_items
+
+
+def to_class_index_mask(
+    annotation: Annotation,
+    label2id: dict,
+    mask_path: str,
+):
+    """Convert annotation to class index mask and save as PNG."""
+    height, width = annotation.img_size
+    class_mask = np.zeros((height, width), dtype=np.uint8)
+
+    for label in annotation.labels:
+        class_name = label.obj_class.name.partition(". ")[2]
+        if class_name not in label2id:
+            tqdm.write(f"Skipping unrecognized label: {label}")
+            continue  # skip unrecognized labels
+
+        class_index = label2id[class_name]
+
+        if label.geometry.geometry_name() == "bitmap":
+            origin = label.geometry.origin
+            top = origin.row
+            left = origin.col
+            bitmap = label.geometry.data  # binary numpy array, shape (h, w)
+
+            h, w = bitmap.shape
+            if top + h > height or left + w > width:
+                tqdm.write(f"Skipping label '{class_name}': size mismatch.")
+                continue
+
+            class_mask[top : top + h, left : left + w][bitmap] = class_index
+        else:
+            continue
+
+    Image.fromarray(class_mask).save(mask_path)
+
+
+def process_datasets(
+    project,
+    datasets,
+    output_dir,
+    label2id,
+    train_items,
+):
+    """Process all datasets and create train/validation splits."""
+    for dataset in tqdm(datasets, desc="Processing datasets"):
+        items = dataset.get_items_names()
+        random.shuffle(items)
+
+        for i, item in tqdm(
+            enumerate(items),
+            desc=f"Processing dataset: {dataset.name}",
+            total=len(items),
+            leave=False,
+        ):
+            # Determine split
+            split = "training" if i < train_items else "validation"
+
+            # Copy images
+            item_paths = dataset.get_item_paths(item)
+            img_path = item_paths.img_path
+            img_filename = os.path.basename(img_path)
+            dest_path = f"{output_dir}/images/{split}/{img_filename}"
+            shutil.copy(img_path, dest_path)
+
+            # Convert and copy annotations
+            ann_path = item_paths.ann_path
+            ann = sly.Annotation.load_json_file(ann_path, project.meta)
+            mask_filename = f"{os.path.splitext(item)[0]}.png"
+            mask_path = f"{output_dir}/annotations/{split}/{mask_filename}"
+            to_class_index_mask(ann, label2id, mask_path)
+
+
+def parse_arguments():
+    """Parse command line arguments."""
+    parser = argparse.ArgumentParser(
+        description="Parse Supervisely project and convert to training format"
+    )
+    parser.add_argument(
+        "--project_dir",
+        type=str,
+        required=True,
+        help="Path to the Supervisely project directory",
+    )
+    parser.add_argument(
+        "--output_base_dir",
+        type=str,
+        required=True,
+        help="Base output directory for parsed data",
+    )
+    parser.add_argument(
+        "--train_ratio",
+        type=float,
+        default=0.8,
+        help="Ratio of data to use for training (default: 0.8)",
+    )
+    parser.add_argument(
+        "--seed",
+        type=int,
+        default=42,
+        help="Random seed for reproducible splits (default: 42)",
+    )
+    return parser.parse_args()
+
+
+def main():
+    """Main function to parse Supervisely project."""
+    # Parse arguments
+    args = parse_arguments()
+
+    # Set random seed for reproducible splits
+    random.seed(args.seed)
+
+    # Load project
+    project = sly.Project(args.project_dir, sly.OpenMode.READ)
+    print(f"Project: {project.name}")
+
+    # Setup output directory
+    output_dir = os.path.join(args.output_base_dir, project.name)
+    create_output_directories(output_dir)
+
+    # Extract class information
+    id2label, id2color, label2id = extract_class_info(project, output_dir)
+
+    # Get datasets and calculate splits
+    datasets = project.datasets
+    print(f"Datasets: {len(datasets)}")
+    train_items, val_items = calculate_split_counts(datasets, args.train_ratio)
+
+    # Process datasets
+    process_datasets(
+        project,
+        datasets,
+        output_dir,
+        label2id,
+        train_items,
+    )
+
+    print("Processing completed successfully!")
+
+
+if __name__ == "__main__":
+    main()

utils/__init__.py

@@ -0,0 +1,7 @@
+from .parse import run_supervisely_parser
+from .train import train_model
+
+__all__ = [
+    "run_supervisely_parser",
+    "train_model",
+]

utils/parse.py

@@ -0,0 +1,101 @@
+import os
+import tempfile
+import zipfile
+import subprocess
+from pathlib import Path
+
+
+def run_supervisely_parser(
+    project_path: str,
+    train_ratio: float,
+    seed: int,
+) -> str:
+    """Extract a Supervisely project zip and run the parser script inside .venv-sly.
+
+    Parameters
+    ----------
+    project_path : (str)
+        Path to the uploaded Supervisely project .zip.
+    train_ratio : float
+        Portion of data to allocate to training (remainder is validation).
+    seed : int
+        Random seed forwarded to the parser for reproducible splits.
+
+    Returns
+    -------
+    str
+        Path to the parsed dataset directory produced by the parser script.
+    """
+
+    project_zip = Path(project_path)
+    if not project_zip.exists():
+        raise FileNotFoundError(
+            f"Provided project zip not found: {project_zip}"
+        )
+    if project_zip.suffix.lower() != ".zip":
+        raise ValueError("Supervisely project must be a .zip archive")
+
+    project_dir = project_zip.parent
+    extract_dir = Path(tempfile.mkdtemp(dir=project_dir))
+    output_base_dir = Path(tempfile.mkdtemp(dir=project_dir))
+
+    with zipfile.ZipFile(project_zip, "r") as zf:
+        zf.extractall(extract_dir)
+
+    def find_project_root(root: Path) -> Path:
+        if (root / "meta.json").exists():
+            return root
+        for child in root.iterdir():
+            if child.is_dir() and (child / "meta.json").exists():
+                return child
+        raise FileNotFoundError(
+            f"Could not locate 'meta.json' inside extracted archive at {root}"
+        )
+
+    project_root = find_project_root(extract_dir)
+
+    repo_root = Path(__file__).resolve().parent.parent
+    parser_script = repo_root / "scripts" / "supervisely_parser.py"
+    venv_python = repo_root / ".venv-sly" / "bin" / "python"
+
+    if not parser_script.exists():
+        raise FileNotFoundError(
+            f"Parser script not found: {parser_script}",
+        )
+    if not venv_python.exists():
+        raise FileNotFoundError(
+            "Expected .venv-sly Python interpreter at: " f"{venv_python}",
+        )
+
+    cmd = [
+        str(venv_python),
+        str(parser_script),
+        "--project_dir",
+        str(project_root),
+        "--output_base_dir",
+        str(output_base_dir),
+        "--train_ratio",
+        str(train_ratio),
+        "--seed",
+        str(seed),
+    ]
+
+    result = subprocess.run(
+        cmd,
+        capture_output=True,
+        text=True,
+        env={**os.environ},
+    )
+    if result.returncode != 0:
+        raise RuntimeError(
+            "Supervisely parser failed.\n"
+            f"STDOUT:\n{result.stdout}\nSTDERR:\n{result.stderr}"
+        )
+    produced_dirs = [p for p in output_base_dir.iterdir() if p.is_dir()]
+    if len(produced_dirs) != 1:
+        raise RuntimeError(
+            "Could not unambiguously determine parsed dataset directory in "
+            f"{output_base_dir}. Found: {produced_dirs}"
+        )
+    dataset_dir = produced_dirs[0]
+    return str(dataset_dir)

utils/train.py

@@ -0,0 +1,602 @@
+"""
+SegFormer Fine-tuning Script
+
+This script fine-tunes a SegFormer model on a custom semantic segmentation
+dataset. It provides configurable parameters for training hyperparameters
+and dataset settings.
+"""
+
+import json
+import os
+import zipfile
+import numpy as np
+import torch
+import torch.nn as nn
+from torch.utils.data import Dataset, DataLoader
+from PIL import Image
+from pathlib import Path
+from datetime import datetime
+from transformers import (
+    SegformerImageProcessor,
+    SegformerForSemanticSegmentation,
+)
+import evaluate
+from tqdm import tqdm
+
+
+class SemanticSegmentationDataset(Dataset):
+    """Image (semantic) segmentation dataset."""
+
+    def __init__(
+        self,
+        root_dir,
+        image_processor,
+        train=True,
+        data_percent=100,
+    ):
+        """
+        Args:
+            root_dir (string): Root directory of the dataset containing
+                the images + annotations.
+            image_processor (SegFormerImageProcessor): image processor to
+                prepare images + segmentation maps.
+            train (bool): Whether to load "training" or "validation"
+                images + annotations.
+            data_percent (int): Percentage of the dataset to use.
+                100 means all data, 50 means half of the data.
+        """
+        self.root_dir = root_dir
+        self.image_processor = image_processor
+        self.train = train
+
+        sub_path = "training" if self.train else "validation"
+        self.img_dir = os.path.join(self.root_dir, "images", sub_path)
+        self.ann_dir = os.path.join(self.root_dir, "annotations", sub_path)
+
+        # read images
+        image_file_names = []
+        for root, dirs, files in os.walk(self.img_dir):
+            image_file_names.extend(files)
+        self.images = sorted(image_file_names)
+
+        # read annotations
+        annotation_file_names = []
+        for root, dirs, files in os.walk(self.ann_dir):
+            annotation_file_names.extend(files)
+        self.annotations = sorted(annotation_file_names)
+
+        assert len(self.images) == len(
+            self.annotations
+        ), "There must be as many images as there are segmentation maps"
+
+        # Apply data_percent to limit the dataset size
+        data_percent = data_percent / 100.0
+        if data_percent < 1.0:
+            images_num_samples = int(len(self.images) * data_percent)
+            annotations_num_samples = int(len(self.annotations) * data_percent)
+            self.images = self.images[:images_num_samples]
+            self.annotations = self.annotations[:annotations_num_samples]
+
+    def __len__(self):
+        return len(self.images)
+
+    def __getitem__(self, idx):
+        image = Image.open(os.path.join(self.img_dir, self.images[idx]))
+        segmentation_map = Image.open(
+            os.path.join(
+                self.ann_dir,
+                self.annotations[idx],
+            ),
+        )
+        encoded_inputs = self.image_processor(
+            image,
+            segmentation_map,
+            return_tensors="pt",
+        )
+
+        for k, v in encoded_inputs.items():
+            encoded_inputs[k].squeeze_()  # remove batch dimension
+
+        return encoded_inputs
+
+
+class MeanDice:
+    def __init__(self):
+        self.reset()
+
+    def reset(self):
+        """Reset stored predictions and references."""
+        self.predictions = []
+        self.references = []
+
+    def add_batch(self, predictions, references):
+        """
+        Add a batch of predictions and references.
+
+        Args:
+            predictions (np.ndarray): Predicted class indices
+            references (np.ndarray): Ground truth class indices
+        """
+        self.predictions.append(predictions)
+        self.references.append(references)
+
+    def compute(self, num_labels, ignore_index=None):
+        """Compute mean Dice score across all stored batches."""
+        predictions = np.concatenate([p.flatten() for p in self.predictions])
+        references = np.concatenate([r.flatten() for r in self.references])
+
+        dice_scores = []
+
+        for class_id in range(num_labels):
+            pred_mask = predictions == class_id
+            ref_mask = references == class_id
+
+            # Exclude ignore_index
+            if ignore_index is not None:
+                valid_mask = references != ignore_index
+                pred_mask = pred_mask & valid_mask
+                ref_mask = ref_mask & valid_mask
+
+            intersection = np.sum(pred_mask & ref_mask)
+            union = np.sum(pred_mask) + np.sum(ref_mask)
+
+            if union == 0:
+                dice = 1.0 if intersection == 0 else 0.0
+            else:
+                dice = 2.0 * intersection / union
+
+            dice_scores.append(dice)
+
+        return {
+            "mean_dice": float(np.mean(dice_scores)),
+            "per_class_dice": dice_scores,
+        }
+
+
+def get_latest_model_dir(base_path: str = "./segformer_finetuned") -> Path:
+    """
+    Returns the Path to the latest model directory based on
+    timestamp folder names.
+
+    Folder names must follow the format: YYYY-MM-DD_HH-MM-SS
+    """
+    base = Path(base_path)
+    if not base.exists() or not base.is_dir():
+        raise FileNotFoundError(f"Directory not found: {base_path}")
+
+    model_dirs = []
+    for d in base.iterdir():
+        if d.is_dir():
+            try:
+                dt = datetime.strptime(d.name, "%Y-%m-%d_%H-%M-%S")
+                model_dirs.append((dt, d))
+            except ValueError:
+                continue  # Skip non-matching directories
+
+    if not model_dirs:
+        raise FileNotFoundError(
+            "No model directories found with valid timestamp format."
+        )
+
+    # Return the directory with the latest timestamp
+    return max(model_dirs, key=lambda x: x[0])[1]
+
+
+def load_model_and_labels(data_dir, model_path):
+    """Load the model and label mappings."""
+    # Load id2label mapping from JSON file
+    id2label = json.load(open(f"{data_dir}/id2label.json", mode="r"))
+    id2label = {int(k): v for k, v in id2label.items()}
+    label2id = {v: k for k, v in id2label.items()}
+
+    # Load id2color mapping from JSON file
+    id2color = json.load(open(f"{data_dir}/id2color.json", "r"))
+
+    print(f"Loaded {len(id2label)} classes:")
+    for i, label in id2label.items():
+        print(f"  {i}: {label}")
+
+    # Load model
+    model = SegformerForSemanticSegmentation.from_pretrained(
+        model_path,
+        num_labels=len(id2label),
+        id2label=id2label,
+        label2id=label2id,
+    )
+    return model, id2label, id2color
+
+
+def create_datasets_and_dataloaders(
+    image_width,
+    image_height,
+    data_dir,
+    batch_size,
+    data_percent,
+):
+    """Create datasets and dataloaders."""
+    image_processor = SegformerImageProcessor(
+        size={"height": image_height, "width": image_width},
+    )
+
+    train_dataset = SemanticSegmentationDataset(
+        root_dir=data_dir,
+        image_processor=image_processor,
+        train=True,
+        data_percent=data_percent,
+    )
+
+    valid_dataset = SemanticSegmentationDataset(
+        root_dir=data_dir,
+        image_processor=image_processor,
+        train=False,
+        data_percent=data_percent,
+    )
+
+    print(f"Number of training examples: {len(train_dataset)}")
+    print(f"Number of validation examples: {len(valid_dataset)}")
+
+    train_dataloader = DataLoader(
+        train_dataset,
+        batch_size=batch_size,
+        shuffle=True,
+    )
+    valid_dataloader = DataLoader(
+        valid_dataset,
+        batch_size=batch_size,
+    )
+
+    return train_dataloader, valid_dataloader
+
+
+def class_indices_to_rgb(class_indices, id2color):
+    """Convert class indices to RGB colored image."""
+    # class_indices shape: (H, W) with integer class IDs
+    height, width = class_indices.shape
+    rgb_image = np.zeros((height, width, 3), dtype=np.uint8)
+
+    for class_id, color in id2color.items():
+        rgb_image[class_indices == class_id] = color
+
+    return rgb_image
+
+
+def validate_model(
+    model: SegformerForSemanticSegmentation,
+    dataloader,
+    device,
+    id2label,
+    calc_dice=False,
+    epoch=None,
+):
+    """
+    Validate the model on a validation set and return loss, IoU, accuracy.
+    """
+    model.eval()
+    metric = evaluate.load("mean_iou")
+    dice = MeanDice()
+    total_loss = 0.0
+    num_batches = 0
+
+    with torch.no_grad():
+        for batch in tqdm(
+            dataloader,
+            desc="Validating Epoch " + str(epoch if epoch is not None else ""),
+            leave=False,
+            unit="batches",
+        ):
+            pixel_values = batch["pixel_values"].to(device)
+            labels = batch["labels"].to(device)
+
+            outputs = model(pixel_values=pixel_values, labels=labels)
+            logits = outputs.logits
+            loss = outputs.loss
+
+            total_loss += loss.item()
+            num_batches += 1
+
+            upsampled_logits = nn.functional.interpolate(
+                logits,
+                size=labels.shape[-2:],
+                mode="bilinear",
+                align_corners=False,
+            )
+            predicted = upsampled_logits.argmax(dim=1)
+
+            # Store predictions and references for additional metrics
+            pred_np = predicted.detach().cpu().numpy()
+            ref_np = labels.detach().cpu().numpy()
+
+            metric.add_batch(
+                predictions=pred_np,
+                references=ref_np,
+            )
+            if calc_dice:
+                dice.add_batch(
+                    predictions=pred_np,
+                    references=ref_np,
+                )
+
+    # Calculate IoU and accuracy
+    result = metric.compute(
+        num_labels=len(id2label),
+        ignore_index=10,
+        reduce_labels=False,
+    )
+    if calc_dice:
+        dice_result = dice.compute(
+            num_labels=len(id2label),
+            ignore_index=10,
+        )
+
+    avg_loss = total_loss / num_batches if num_batches > 0 else 0.0
+    return (
+        avg_loss,
+        result["mean_iou"],
+        result["per_category_iou"],
+        result["mean_accuracy"],
+        result["per_category_accuracy"],
+        dice_result["mean_dice"] if calc_dice else None,
+        dice_result["per_class_dice"] if calc_dice else None,
+    )
+
+
+def run_training(
+    model: SegformerForSemanticSegmentation,
+    device,
+    train_dataloader,
+    valid_dataloader,
+    id2label,
+    num_epochs,
+    learning_rate,
+    early_stopping,
+    validate_every,
+):
+    """Train the model.
+
+    Returns
+    -------
+    tuple(best_model, metrics)
+        best_model : nn.Module
+        metrics : dict with lists for keys: 'epoch', 'train_loss', 'train_iou',
+          'train_acc', 'val_loss', 'val_iou', 'val_acc'
+    """
+    # Setup device
+    model.to(device)
+
+    # Setup optimizer
+    optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate)
+
+    # Setup metrics
+    metrics = {
+        "epoch": [],
+        "train_loss": [],
+        "train_iou": [],
+        "train_acc": [],
+        "val_loss": [],
+        "val_iou": [],
+        "val_acc": [],
+    }
+
+    metric = evaluate.load("mean_iou")
+
+    model.train()
+
+    # Initial validation
+    (
+        loss,
+        iou,
+        per_class_iou,
+        acc,
+        per_class_acc,
+        dice,
+        dice_per_class,
+    ) = validate_model(
+        model=model,
+        dataloader=valid_dataloader,
+        device=device,
+        id2label=id2label,
+        calc_dice=True,
+        epoch=0,
+    )
+    # Add to metrics at epoch 0
+    metrics["epoch"].append(int(0))
+    metrics["val_loss"].append(loss)
+    metrics["val_iou"].append(iou)
+    metrics["val_acc"].append(acc)
+    metrics["train_loss"].append(None)
+    metrics["train_iou"].append(None)
+    metrics["train_acc"].append(None)
+
+    initial_dice = dice
+
+    best_model = model
+
+    best_iou = iou
+    patience = early_stopping
+    epochs_without_improvement = 0
+    for epoch in tqdm(
+        range(num_epochs),
+        desc="Training Epochs",
+        unit="epochs",
+    ):
+        epoch_loss = 0.0
+        num_batches = 0
+        model.train()  # Ensure model is in training mode
+
+        progress_bar = tqdm(
+            train_dataloader,
+            desc=f"Training Epoch {epoch + 1}",
+            leave=True,
+            unit="batches",
+        )
+
+        for idx, batch in enumerate(progress_bar):
+            # Get the inputs
+            pixel_values = batch["pixel_values"].to(device)
+            labels = batch["labels"].to(device)
+
+            # Zero the parameter gradients
+            optimizer.zero_grad()
+
+            # Forward + backward + optimize
+            outputs = model(pixel_values=pixel_values, labels=labels)
+            loss, logits = outputs.loss, outputs.logits
+
+            loss.backward()
+            optimizer.step()
+
+            epoch_loss += loss.item()
+            num_batches += 1
+
+            # Evaluate training batch
+            with torch.no_grad():
+                upsampled_logits = nn.functional.interpolate(
+                    logits,
+                    size=labels.shape[-2:],
+                    mode="bilinear",
+                    align_corners=False,
+                )
+                predicted = upsampled_logits.argmax(dim=1)
+
+                # Store for metric calculation
+                pred_np = predicted.detach().cpu().numpy()
+                ref_np = labels.detach().cpu().numpy()
+
+                # Note: metric expects predictions + labels as numpy arrays
+                metric.add_batch(
+                    predictions=pred_np,
+                    references=ref_np,
+                )
+
+        train_metrics = metric.compute(
+            num_labels=len(id2label),
+            ignore_index=10,
+            reduce_labels=False,
+        )
+        train_loss = epoch_loss / num_batches if num_batches else 0.0
+
+        # Validation
+        if (epoch + 1) % validate_every == 0:
+            (
+                val_loss,
+                val_iou,
+                val_per_class_iou,
+                val_acc,
+                val_per_class_acc,
+                val_dice,
+                val_dice_per_class,
+            ) = validate_model(
+                model=model,
+                dataloader=valid_dataloader,
+                device=device,
+                id2label=id2label,
+                epoch=epoch + 1,
+            )
+
+        # Record metrics
+        metrics["epoch"].append(int(epoch + 1))
+        metrics["train_loss"].append(train_loss)
+        metrics["train_iou"].append(train_metrics["mean_iou"])
+        metrics["train_acc"].append(train_metrics["mean_accuracy"])
+        metrics["val_loss"].append(val_loss)
+        metrics["val_iou"].append(val_iou)
+        metrics["val_acc"].append(val_acc)
+
+        # Save the best model
+        if val_iou > best_iou:
+            best_model = model
+            best_iou = val_iou
+            epochs_without_improvement = 0
+        else:
+            epochs_without_improvement += 1
+
+        if epochs_without_improvement >= patience:
+            tqdm.write(
+                f"Early stopping after {patience} epochs with no improvement",
+            )
+            break
+
+    return best_model, metrics, initial_dice
+
+
+def extract_model_zip(model_zip_path):
+    """Extract model zip file and return the model directory."""
+
+    if not os.path.exists(model_zip_path):
+        raise FileNotFoundError(f"Model zip file not found: {model_zip_path}")
+
+    with zipfile.ZipFile(model_zip_path, "r") as zip_ref:
+        extract_dir = os.path.join(os.path.dirname(model_zip_path), "output")
+        zip_ref.extractall(extract_dir)
+
+    # Check nested folder
+    if len(os.listdir(extract_dir)) == 1:
+        return os.path.join(extract_dir, os.listdir(extract_dir)[0])
+    else:
+        return extract_dir
+
+
+def train_model(
+    data_dir,
+    base_model_zip,
+    image_width,
+    image_height,
+    batch_size,
+    data_percent,
+    num_epochs,
+    learning_rate,
+    early_stopping,
+    validate_every,
+):
+
+    model_path = extract_model_zip(base_model_zip)
+
+    # Load model and labels
+    model, id2label, id2color = load_model_and_labels(data_dir, model_path)
+
+    # Create datasets and dataloaders
+    train_dataloader, valid_dataloader = create_datasets_and_dataloaders(
+        image_width,
+        image_height,
+        data_dir,
+        batch_size,
+        data_percent,
+    )
+
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    print(f"Using device: {device}")
+
+    # Train the model
+    best_model, metrics, initial_dice = run_training(
+        model,
+        device,
+        train_dataloader,
+        valid_dataloader,
+        id2label,
+        num_epochs,
+        learning_rate,
+        early_stopping,
+        validate_every,
+    )
+
+    # Final validation
+    (
+        loss,
+        iou,
+        per_class_iou,
+        acc,
+        per_class_acc,
+        dice,
+        dice_per_class,
+    ) = validate_model(
+        model=best_model,
+        dataloader=valid_dataloader,
+        device=device,
+        id2label=id2label,
+        calc_dice=True,
+        epoch=0,
+    )
+
+    final_dice = dice
+
+    return best_model, metrics, [initial_dice, final_dice]