add metrics computation

app.py

@@ -2,6 +2,7 @@ import os
 
 import gradio as gr
 import numpy as np
+import pandas as pd
 import torch
 from patchify import patchify, unpatchify
 from phasepack import phasecong
@@ -12,6 +13,8 @@ from skimage.feature import canny
 from skimage.filters import sato
 
 from src.unet import UNet
+from src.train import eval_single
+from src.dataset_benchm import expand_wide_fractures_gt, dilate_labels
 
 # ------------------------------------------------------------
 # Device
@@ -103,6 +106,31 @@ def sato_fn(img, x, sigmas):
     return np.float64(sato(gray, sato_sigmas_list[sigmas]) < x)
 
 
+# ------------------------------------------------------------
+# Compute metrics
+# ------------------------------------------------------------
+def compute_metrics_ui(gt_img, pred_img, threshold):
+    if gt_img is None or pred_img is None:
+        return None
+
+    # Normalise to [0,1]
+    gt = np.array(gt_img, dtype=np.uint8)
+    pred = np.array(pred_img, dtype=np.uint8)
+
+    if gt.ndim == 3:
+        gt = gt[..., 0]
+    if pred.ndim == 3:
+        pred = pred[..., 0]
+
+    gt = dilate_labels(gt)
+
+    metrics = eval_single(gt, pred, threshold=threshold, device=device)
+
+    df = pd.DataFrame([metrics])
+    df = df.round(3)
+    return df
+
+
 # ------------------------------------------------------------
 # Deep learning model loading
 # ------------------------------------------------------------
@@ -309,7 +337,7 @@ with gr.Blocks(title="Fractex2D Segmentation") as demo:
             gr.Markdown(
                 """
                 ## Canny edge detection
-                Canny edge detection (scikit-image) with normalized thresholds https://doi.org/10.1109/TPAMI.1986.4767851.
+                Canny edge detection (scikit-image) with normalised thresholds https://doi.org/10.1109/TPAMI.1986.4767851.
                 - **sigma** controls Gaussian smoothing
                 - **lt / ht** are low/high thresholds in the range 0–1
                 """
@@ -395,6 +423,59 @@ with gr.Blocks(title="Fractex2D Segmentation") as demo:
                 outputs=pc_out,
             )
 
+        # ------------------------------------------------------------
+        # TAB 5 — METRICS
+        # ------------------------------------------------------------
+        with gr.Tab("Metrics computation"):
+            gr.Markdown(
+                """
+                ## Segmentation Metrics
+                Compute quantitative metrics between a **prediction** and a **ground-truth** (1px wide annotation).
+                Both images must be aligned and have the same resolution.
+                """
+            )
+
+            with gr.Row():
+                gt_input = gr.Image(label="Ground truth", type="numpy")
+                pred_input = gr.Image(label="Prediction", type="numpy")
+
+            with gr.Row():
+                thresh = gr.Slider(
+                    0, 1,
+                    value=0.1,
+                    step=0.01,
+                    label="Binarisation threshold"
+                )
+
+            with gr.Row():
+                with gr.Column(scale=1):
+                    pass
+                metric_btn = gr.Button("Compute metrics")
+                with gr.Column(scale=1):
+                    pass
+
+            metric_table = gr.Dataframe(
+                headers=[
+                    "mse", "psnr", "ssim", "ae",
+                    "acc", "prec", "rec", "spec",
+                    "f1", "dice", "iou", "ck", "roc_auc"
+                ],
+                label="Metrics (single image pair)"
+            )
+
+            metric_btn.click(
+                fn=compute_metrics_ui,
+                inputs=[gt_input, pred_input, thresh],
+                outputs=metric_table,
+            )
+
+            gr.Examples(
+                examples=[
+                    ["examples/kl5-s3_1-gt.png", "examples/unet-p1_pred_kl5-s3_1.png", 0.1],
+                ],
+                inputs=[gt_input, pred_input, thresh],
+            )
+
     # ------------------------------------------------------------
     # Extra reference
     # ------------------------------------------------------------

src/dataset_benchm.py

@@ -0,0 +1,524 @@
+from pathlib import Path
+import random
+from typing import List, Optional, Sequence, Tuple
+
+import numpy as np
+import torch
+import torchvision.transforms.v2 as t
+import torchvision.transforms.v2.functional as TF
+from skimage import io
+from skimage.filters.rank import maximum
+from skimage.measure import label
+from skimage.morphology import binary_dilation, dilation, disk
+from skimage.segmentation import expand_labels
+from torch.utils.data import ConcatDataset, DataLoader, Dataset
+
+
+# -------------------------
+# Label pre-processing
+# -------------------------
+def expand_wide_fractures_gt(
+    img: np.ndarray,
+    gt: np.ndarray,
+    disk_size: int = 2,
+    thresh: int = 30,
+    gt_thresh: int = 100,
+    gt_ext: str = "png",
+) -> np.ndarray:
+    """
+    Expand a binary/soft ground-truth mask to include nearby wide/dark fractures.
+
+    Method:
+      - Use green channel (index 1) as a grayscale proxy.
+      - Apply a maximum filter to emphasize large dark regions.
+      - Threshold and dilate to form a candidate mask.
+      - Keep only connected components that overlap the original GT.
+      - Return a combined mask as uint8 (0..255). If gt_ext contains "tif" the
+        original `gt` is assumed to be already in [0,1] or in the original dtype;
+        the code preserves existing scaling behavior from the original script.
+
+    Args:
+        img: HxWxC image (expects at least 2 channels; green channel used).
+        gt: HxW ground-truth mask (expected in [0..1] or [0..255]).
+        disk_size: radius for morphological operations.
+        thresh: threshold applied to the maximum-filtered gray image.
+        gt_thresh: threshold to consider a pixel part of the original GT.
+        gt_ext: file extension of GT (affects final combination step).
+
+    Returns:
+        Expanded GT mask as np.uint8 (values 0 or 255).
+    """
+    if img.ndim < 3 or img.shape[2] < 2:
+        raise ValueError("img must have at least 2 channels (uses green channel).")
+
+    # use green channel as grayscale proxy
+    gray = img[..., 1].astype(np.uint8)
+
+    # keep large dark areas via maximum filter, then threshold and dilate
+    imax = maximum(gray, disk(disk_size))
+    candidate = binary_dilation(imax < thresh, disk(disk_size))
+
+    # combine candidate with existing GT (considering gt_thresh)
+    gt_bool = gt > gt_thresh
+    combined = np.logical_or(candidate, gt_bool)
+
+    # remove connected components that do not overlap original GT
+    labeled, num = label(combined, connectivity=1, return_num=True)
+    for comp_id in range(1, num + 1):
+        comp_mask = labeled == comp_id
+        if not np.any(gt_bool[comp_mask]):
+            combined[comp_mask] = False
+
+    # produce uint8 [0,255] result with behavior matching original code
+    if "tif" in gt_ext:
+        # preserve original gt scaling behavior from source
+        new_gt = (np.array(gt * 255, dtype=np.uint8) | np.array(combined * 255, dtype=np.uint8))
+    else:
+        new_gt = (np.array(gt, dtype=np.uint8) | np.array(combined * 255, dtype=np.uint8))
+
+    return new_gt
+
+
+def dilate_labels(image: np.ndarray) -> np.ndarray:
+    """
+    Smooth label boundaries by multi-scale dilation and blending.
+
+    - Expand labels to fill tiny gaps (expand_labels).
+    - Create three dilation masks with increasing disks and blend them into
+      a smoothed label map with decreasing weights.
+
+    Args:
+        image: integer-labeled image or binary mask (HxW).
+
+    Returns:
+        np.uint8 array (HxW) with blended/smoothed label boundaries.
+    """
+    expanded = expand_labels(image, distance=2)
+
+    # Multi-scale dilation masks (exclusive differences)
+    d1 = dilation(expanded, disk(2)) ^ expanded
+    d2 = dilation(expanded, disk(5)) ^ d1 ^ expanded
+    d3 = dilation(expanded, disk(7)) ^ d2 ^ d1 ^ expanded
+
+    blended = expanded + d1 / 3.0 + d2 / 5.0 + d3 / 9.0
+    return np.array(blended, dtype=np.uint8)
+
+
+# -------------------------
+# Augmentation helpers
+# -------------------------
+def _apply_random_flips(image: torch.Tensor, mask: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+    """Random horizontal and vertical flips (50% each)."""
+    if random.random() > 0.5:
+        image, mask = TF.hflip(image), TF.hflip(mask)
+    if random.random() > 0.5:
+        image, mask = TF.vflip(image), TF.vflip(mask)
+    return image, mask
+
+
+def _apply_random_photometric_augmentations(image: torch.Tensor, prob_config: Optional[dict] = None) -> torch.Tensor:
+    """
+    Photometric augmentations applied independently with small probabilities.
+
+    The function preserves an extra channel (e.g. DEM) if image has 4 channels:
+      - augment only the first three (RGB) channels, then concatenate the extra.
+    """
+    if prob_config is None:
+        prob_config = {
+            "gaussian_blur": 0.05,
+            "darken_low": 0.05,
+            "brighten": 0.15,
+            "contrast": 0.05,
+            "saturation": 0.05,
+        }
+
+    has_extra = image.shape[0] == 4
+    rgb = image[:3] if has_extra else image
+
+    # gaussian blur
+    if random.random() < prob_config["gaussian_blur"]:
+        sigma = random.uniform(0.1, 2.0)
+        rgb = TF.gaussian_blur(rgb, kernel_size=5, sigma=sigma)
+
+    # darken (factor < 1)
+    if random.random() < prob_config["darken_low"]:
+        factor = random.uniform(0.7, 0.9)
+        rgb = TF.adjust_brightness(rgb, factor)
+
+    # brighten (factor > 1)
+    if random.random() < prob_config["brighten"]:
+        factor = random.uniform(1.1, 1.7)
+        rgb = TF.adjust_brightness(rgb, factor)
+
+    # contrast
+    if random.random() < prob_config["contrast"]:
+        factor = random.uniform(0.7, 1.5)
+        rgb = TF.adjust_contrast(rgb, factor)
+
+    # saturation
+    if random.random() < prob_config["saturation"]:
+        factor = random.uniform(0.7, 1.5)
+        rgb = TF.adjust_saturation(rgb, factor)
+
+    if has_extra:
+        image = torch.cat([rgb, image[3:]], dim=0)
+    else:
+        image = rgb
+
+    return image
+
+
+# -------------------------
+# Base dataset utilities
+# -------------------------
+def _read_image(path: Path) -> np.ndarray:
+    """Read image with skimage.io and ensure dtype uint8."""
+    arr = io.imread(str(path))
+    # convert floats to uint8 if necessary
+    if arr.dtype != np.uint8:
+        arr = arr.astype(np.uint8)
+    return arr
+
+
+def _read_mask(path: Path) -> np.ndarray:
+    """Read mask and convert to uint8 0..255."""
+    arr = io.imread(str(path))
+    if arr.dtype != np.uint8:
+        arr = (arr * 255).astype(np.uint8) if arr.max() <= 1.0 else arr.astype(np.uint8)
+    return arr
+
+
+# -------------------------
+# Dataset classes
+# -------------------------
+class BaseCrackDataset(Dataset):
+    """
+    Minimal common functionality for the specific dataset wrappers used downstream.
+
+    Subclasses must provide:
+      - self.images (list[Path])
+      - self.masks  (list[Path])
+      - optional self.dems  (list[Path]) when in_channels==4
+    """
+
+    def __init__(
+        self,
+        images: Sequence[Path],
+        masks: Sequence[Path],
+        dem_paths: Optional[Sequence[Path]] = None,
+        topo: bool = False,
+        transform: bool = False,
+        expand: bool = True,
+        dilate: bool = True,
+        in_channels: int = 3,
+    ):
+        self.images = list(images)
+        self.masks = list(masks)
+        self.dems = list(dem_paths) if dem_paths is not None else None
+
+        self.topo = topo
+        self.transform = transform
+        self.expand = expand
+        self.dilate = dilate
+        self.in_channels = in_channels
+
+    def __len__(self) -> int:
+        return len(self.images)
+
+    def _load_pair(self, idx: int) -> Tuple[torch.Tensor, torch.Tensor]:
+        """
+        Load image/mask pair, apply optional expand/dilate and channel handling,
+        then perform flips and photometric augmentations.
+        """
+        img_np = _read_image(Path(self.images[idx]))
+        gt_np = _read_mask(Path(self.masks[idx]))
+
+        # expand wide fractures (if requested)
+        if self.expand:
+            gt_np = expand_wide_fractures_gt(img_np[:, :, :3].astype(np.uint8), gt_np)
+
+        # dilate labels (if requested)
+        if self.dilate:
+            gt_np = dilate_labels(gt_np)
+
+        # build image tensor. If dataset provides DEM as a separate file, append as 4th channel.
+        img_tensor = torch.from_numpy(img_np[:, :, :3])
+        if self.in_channels == 4:
+            # if DEM present inside the image array or as separate file, handle both cases
+            if img_np.shape[2] >= 4:
+                dem_np = img_np[:, :, 3].astype(np.float32)
+            elif self.dems is not None:
+                dem_np = _read_image(Path(self.dems[idx])).astype(np.float32)
+            else:
+                raise RuntimeError("Requested 4 input channels but no DEM found.")
+            # normalize DEM to [0,1]
+            dem_tensor = torch.from_numpy(dem_np).float()
+            dem_tensor = (dem_tensor - dem_tensor.min()) / (dem_tensor.max() - dem_tensor.min() + 1e-8)
+            img_tensor = torch.cat((img_tensor, dem_tensor.unsqueeze(2)), axis=2)
+
+        # reformat to C,H,W and normalize image to [0,1]
+        img_tensor = img_tensor.permute(2, 0, 1).float() / 255.0
+
+        mask_tensor = torch.from_numpy(gt_np).unsqueeze(0).float() / 255.0
+
+        # random flips
+        img_tensor, mask_tensor = _apply_random_flips(img_tensor, mask_tensor)
+
+        # photometric augmentations
+        if self.transform:
+            img_tensor = _apply_random_photometric_augmentations(img_tensor)
+
+        return img_tensor.float(), mask_tensor.float()
+
+    def __getitem__(self, index: int) -> Tuple[torch.Tensor, torch.Tensor]:
+        idx = index % len(self.images)
+        return self._load_pair(idx)
+
+
+# -------------------------
+# Concrete dataset wrappers
+# -------------------------
+def _read_list_file(list_path: Path) -> List[str]:
+    """Read non-empty lines from a list file and return them as strings."""
+    with list_path.open("r") as f:
+        return [ln.strip() for ln in f if ln.strip()]
+
+
+class OVAS(BaseCrackDataset):
+    """OVAS dataset wrapper. Expects directory structure: <root>/<subset>/{image,gt,dem}."""
+
+    def __init__(
+        self,
+        subset: str,
+        list_file: Optional[str] = "list.txt",
+        topo: bool = False,
+        transform: bool = False,
+        expand: bool = True,
+        dilate: bool = True,
+        in_channels: int = 3,
+    ):
+        root = Path("data/ovaskainen23_") / subset
+        ext_img = "png"
+        ext_gt = "tif"
+
+        names = []
+        if list_file:
+            names = _read_list_file(root / list_file)
+
+            images = [
+                (root / "image" / n).with_suffix("." + ext_img)
+                for n in names
+                if n.endswith("." + ext_gt)
+            ]
+            masks = [root / "gt" / n for n in names if n.endswith("." + ext_gt)]
+            dems = [root / "dem" / n for n in names if n.endswith("." + ext_gt)]
+        else:
+            images = sorted(path for path in (root / "image").iterdir() if path.suffix.lower().lstrip(".") == ext_img)
+            masks = sorted(path for path in (root / "gt").iterdir() if path.suffix.lower().lstrip(".") == ext_gt)
+            dems = sorted(path for path in (root / "dem").iterdir() if path.suffix.lower().lstrip(".") == ext_gt)
+
+        super().__init__(images=images, masks=masks, dem_paths=dems, topo=topo, transform=transform,
+                         expand=expand, dilate=dilate, in_channels=in_channels)
+
+
+class MATTEO(BaseCrackDataset):
+    """MATTEO dataset wrapper. Expects .tif files; includes DEM channel inside the image."""
+
+    def __init__(
+        self,
+        subset: str,
+        list_file: Optional[str] = "list.txt",
+        topo: bool = False,
+        transform: bool = False,
+        expand: bool = True,
+        dilate: bool = True,
+        in_channels: int = 3,
+    ):
+        root = Path("data/matteo21") / subset
+        ext = "tif"
+
+        if list_file:
+            names = _read_list_file(root / list_file)
+        else:
+            names = [p.name for p in (root / "image").iterdir() if p.suffix.lstrip(".") == ext]
+
+        images = sorted(root / "image" / name for name in names)
+        masks = sorted(root / "gt" / name for name in names)
+
+        super().__init__(images=images, masks=masks, dem_paths=None, topo=topo, transform=transform,
+                         expand=expand, dilate=dilate, in_channels=in_channels)
+
+
+class SAMSU(BaseCrackDataset):
+    """SAMSU dataset wrapper. Similar layout to OVAS."""
+
+    def __init__(
+        self,
+        subset: str,
+        list_file: Optional[str] = "list.txt",
+        topo: bool = False,
+        transform: bool = False,
+        expand: bool = True,
+        dilate: bool = True,
+        in_channels: int = 3,
+    ):
+        root = Path("data/samsu19") / subset
+        ext_img = "png"
+        ext_gt = "tif"
+
+        names = []
+        if list_file:
+            names = _read_list_file(root / list_file)
+            images = [
+                (root / "image" / n).with_suffix("." + ext_img)
+                for n in names
+                if n.endswith("." + ext_gt)
+            ]
+            masks = [root / "gt" / n for n in names if n.endswith("." + ext_gt)]
+            dems = [root / "dem" / n for n in names if n.endswith("." + ext_gt)]
+        else:
+            images = sorted(p for p in (root / "image").iterdir() if p.suffix.lstrip(".") == ext_img)
+            masks = sorted(p for p in (root / "gt").iterdir() if p.suffix.lstrip(".") == ext_gt)
+            dems = sorted(p for p in (root / "dem").iterdir() if p.suffix.lstrip(".") == ext_gt)
+
+        super().__init__(images=images, masks=masks, dem_paths=dems, topo=topo, transform=transform,
+                         expand=expand, dilate=dilate, in_channels=in_channels)
+
+
+class GeoCrack(BaseCrackDataset):
+    """GeoCrack dataset wrapper (simple PNG images)."""
+
+    def __init__(
+        self,
+        subset: str,
+        topo: bool = False,
+        transform: bool = False,
+        expand: bool = True,
+        dilate: bool = True,
+        in_channels: int = 3,
+    ):
+        root = Path("data/GeoCrack_") / subset
+        ext = "png"
+
+        images = sorted(p for p in (root / "image").iterdir() if p.suffix.lstrip(".") == ext)
+        masks = sorted(p for p in (root / "gt").iterdir() if p.suffix.lstrip(".") == ext)
+
+        super().__init__(images=images, masks=masks, dem_paths=None, topo=topo, transform=transform,
+                         expand=expand, dilate=dilate, in_channels=in_channels)
+
+    def __getitem__(self, index: int) -> Tuple[torch.Tensor, torch.Tensor]:
+        img, mask = super().__getitem__(index)
+        # consistent resizing used originally
+        img = t.Resize(256)(img)
+        mask = t.Resize(256)(mask)
+        return img.float(), mask.float()
+
+
+class DIC(BaseCrackDataset):
+    """DIC dataset wrapper: single-channel images and PNG masks."""
+
+    def __init__(
+        self,
+        subset: str,
+        topo: bool = False,
+        transform: bool = False,
+        expand: bool = False,
+        dilate: bool = False,
+        in_channels: int = 1,
+    ):
+        root = Path("data/DIC") / subset
+        ext_img = "tif"
+        ext_mask = "png"
+
+        images = sorted(p for p in (root / "image").iterdir() if p.suffix.lstrip(".") == ext_img)
+        masks = sorted(p for p in (root / "gt").iterdir() if p.suffix.lstrip(".") == ext_mask)
+
+        super().__init__(images=images, masks=masks, dem_paths=None, topo=topo, transform=transform,
+                         expand=expand, dilate=dilate, in_channels=in_channels)
+
+    def _load_pair(self, idx: int) -> Tuple[torch.Tensor, torch.Tensor]:
+        """
+        Override to handle single-channel image format (the base expects >=3 channels).
+        """
+        img_np = _read_image(Path(self.images[idx]))
+        gt_np = _read_mask(Path(self.masks[idx]))
+
+        # ensure single channel
+        if img_np.ndim == 3:
+            img_np = img_np[..., 0]
+
+        img_tensor = torch.from_numpy(img_np).unsqueeze(0).float() / 255.0
+        mask_tensor = torch.from_numpy(gt_np).unsqueeze(0).float() / 255.0
+
+        img_tensor, mask_tensor = _apply_random_flips(img_tensor, mask_tensor)
+
+        if self.transform:
+            img_tensor = _apply_random_photometric_augmentations(img_tensor)
+
+        img_tensor = t.Resize(256)(img_tensor)
+        mask_tensor = t.Resize(256)(mask_tensor)
+
+        return img_tensor.float(), mask_tensor.float()
+
+
+# -------------------------
+# Dataset registry & loader builder
+# -------------------------
+DATASETS = {
+    "ovaskainen23": OVAS,
+    "matteo21": MATTEO,
+    "samsu19": SAMSU,
+    "geocrack": GeoCrack,
+    "dic": DIC,
+}
+
+
+def all_datasets(
+    batch_size: int = 32,
+    datasets: str = "samsu19-matteo21-ovaskainen23",
+    in_channels: int = 4,
+    out_channels: int = 1,
+    shape: int = 256,
+    expand: bool = True,
+    dilate: bool = True,
+    shuffle_train: bool = True,
+    do_transform: bool = True,
+) -> Tuple[DataLoader, DataLoader, DataLoader]:
+    """
+    Create concatenated train/val/test DataLoaders from multiple dataset names.
+
+    Args:
+        batch_size: batch size for DataLoaders.
+        datasets: dash-separated dataset keys from DATASETS dict.
+        in_channels: number of input channels requested (3 or 4).
+        out_channels: number of output channels (kept for API compatibility).
+        shape: target shape (not used directly here; datasets may resize internally).
+        expand, dilate: whether to apply expand/dilate preprocessing.
+        shuffle_train: whether to shuffle the training DataLoader.
+        do_transform: whether to enable augmentations.
+
+    Returns:
+        Tuple(train_loader, val_loader, test_loader)
+    """
+    keys = [k.strip() for k in datasets.split("-") if k.strip()]
+    all_train = []
+    all_val = []
+    all_test = []
+
+    for name in keys:
+        if name not in DATASETS:
+            raise KeyError(f"Unknown dataset key: {name}")
+        DS = DATASETS[name]
+        all_train.append(DS(subset="train", transform=do_transform, expand=expand, dilate=dilate, in_channels=in_channels))
+        all_val.append(DS(subset="valid", transform=False, expand=expand, dilate=dilate, in_channels=in_channels))
+        all_test.append(DS(subset="test", transform=False, expand=expand, dilate=dilate, in_channels=in_channels))
+
+    trainset = ConcatDataset(all_train)
+    valset = ConcatDataset(all_val)
+    testset = ConcatDataset(all_test)
+
+    trainloader = DataLoader(trainset, batch_size=batch_size, shuffle=shuffle_train)
+    valloader = DataLoader(valset, batch_size=batch_size, shuffle=False)
+    testloader = DataLoader(testset, batch_size=batch_size, shuffle=False)
+
+    return trainloader, valloader, testloader

src/train.py

@@ -0,0 +1,222 @@
+import matplotlib.pyplot as plt
+import numpy as np
+import torch
+from skimage.morphology import label, skeletonize
+from skimage.util import view_as_windows
+from torchmetrics import MeanAbsoluteError, MeanSquaredError
+from torchmetrics.classification import (
+    BinaryAccuracy, BinaryAUROC, BinaryCohenKappa, BinaryF1Score,
+    BinaryJaccardIndex, BinaryPrecision, BinaryRecall, BinarySpecificity
+)
+from torchmetrics.image import PeakSignalNoiseRatio, StructuralSimilarityIndexMeasure
+from torchmetrics.segmentation import DiceScore
+from tqdm.auto import tqdm
+
+
+def remove_junctions(skel: np.ndarray) -> np.ndarray:
+    """Remove junction points from a binary skeleton."""
+    skel = skel.astype(np.uint8)
+    mask = np.zeros_like(skel)
+    windows = view_as_windows(skel, (3, 3))
+    for i in range(windows.shape[0]):
+        for j in range(windows.shape[1]):
+            if windows[i, j].sum() > 4:
+                mask[i:i+3, j:j+3] = 1
+    return skel * (1 - mask)
+
+
+def fracture_similarity(pred_mask: torch.Tensor, true_mask: torch.Tensor) -> float:
+    """Compute similarity score between predicted and true fracture masks."""
+    pred_skel = skeletonize((pred_mask > 0.1).cpu().numpy())
+    true_skel = skeletonize((true_mask > 0.1).cpu().numpy())
+    pred_clean = remove_junctions(pred_skel)
+    true_clean = remove_junctions(true_skel)
+    pred_labeled = label(pred_clean)
+    true_labeled = label(true_clean)
+    pred_lengths = np.bincount(pred_labeled.ravel())[1:]
+    true_lengths = np.bincount(true_labeled.ravel())[1:]
+    bins = np.linspace(0, 260, 20)
+    pred_hist, _ = np.histogram(pred_lengths, bins=bins)
+    true_hist, _ = np.histogram(true_lengths, bins=bins)
+    pred_hist = pred_hist + 1e-6
+    true_hist = true_hist + 1e-6
+    chi_dist = 0.5 * np.sum((pred_hist - true_hist)**2 / (pred_hist + true_hist))
+    return chi_dist
+
+
+def train_loop(model, optimizer, criterion, train_loader, device='cpu', mdl=None):
+    """Train the model for one epoch."""
+    running_loss = 0
+    model = model.to(device)
+    model.train()
+    pbar = tqdm(train_loader, desc="Iterating over train data")
+
+    for images, labels in pbar:
+        images, labels = images.to(device), labels.to(device)
+        out = model(images)['out'] if mdl == 'fcn_resnet101' else model(images)
+        loss = criterion(out, labels)
+        running_loss += loss.item() * images.shape[0]
+        optimizer.zero_grad()
+        loss.backward()
+        optimizer.step()
+
+    running_loss /= len(train_loader.sampler)
+    return running_loss
+
+
+def eval_loop(model, scheduler, criterion, eval_loader, threshold=0.5, device='cpu',
+              mdl=None, ignore_index=None):
+    """Evaluate the model on a validation or test dataset."""
+    running_loss = 0
+    model.eval()
+    if ignore_index not in [0, 1]:
+        ignore_index = None
+
+    with torch.no_grad():
+        # Metrics
+        acc_metric = BinaryAccuracy(ignore_index=ignore_index).to(device)
+        f1_metric = BinaryF1Score(ignore_index=ignore_index).to(device)
+        prec_metric = BinaryPrecision(ignore_index=ignore_index).to(device)
+        rec_metric = BinaryRecall(ignore_index=ignore_index).to(device)
+        spec_metric = BinarySpecificity(ignore_index=ignore_index).to(device)
+        auroc_metric = BinaryAUROC(ignore_index=ignore_index).to(device)
+        iou_metric = BinaryJaccardIndex(ignore_index=ignore_index).to(device)
+        dice_metric = DiceScore(num_classes=1, average="micro",
+                                aggregation_level='global').to(device)
+        ck_metric = BinaryCohenKappa().to(device)
+        mse_metric = MeanSquaredError().to(device)
+        ae_metric = MeanAbsoluteError().to(device)
+        psnr_metric = PeakSignalNoiseRatio(data_range=1.0).to(device)
+        ssim_metric = StructuralSimilarityIndexMeasure().to(device)
+        fracture_sim_scores = []
+
+        pbar = tqdm(eval_loader, desc='Iterating over evaluation/test data')
+        for imgs, labels in pbar:
+            imgs, labels = imgs.to(device), labels.to(device)
+            out = model(imgs)['out'] if mdl == 'fcn_resnet101' else model(imgs)
+            loss = criterion(out, labels)
+            running_loss += loss.item() * imgs.shape[0]
+
+            predicted = out
+            if mdl == 'Segformer':
+                predicted[predicted > 0.99] = 0.
+            predicted_clf = (out > threshold).float()
+            labels_clf = (labels > 0.).float()
+            labels = labels.float()
+
+            # Compute metrics
+            acc_metric(predicted_clf, labels_clf)
+            f1_metric(predicted_clf, labels_clf)
+            prec_metric(predicted_clf, labels_clf)
+            rec_metric(predicted_clf, labels_clf)
+            spec_metric(predicted_clf, labels_clf)
+            if labels_clf.numel() > 0 and labels_clf.min() != labels_clf.max():
+                auroc_metric(predicted_clf, labels_clf)
+            dice_metric(predicted_clf, labels_clf)
+            iou_metric(predicted_clf, labels_clf)
+            ck_metric(predicted_clf, labels_clf)
+            mse_metric(predicted, labels)
+            psnr_metric(predicted, labels)
+            ssim_metric(predicted, labels)
+            ae_metric(predicted, labels)
+
+            for i in range(imgs.shape[0]):
+                pred_mask = predicted_clf[i, 0].detach().cpu()
+                true_mask = labels_clf[i, 0].detach().cpu()
+                fracture_sim_scores.append(fracture_similarity(pred_mask, true_mask))
+
+        avg_fracture_sim = float(np.mean(fracture_sim_scores)) if fracture_sim_scores else float('nan')
+
+    return {
+        'mse': mse_metric.compute().item(),
+        'psnr': psnr_metric.compute().item(),
+        'ssim': ssim_metric.compute().item(),
+        'ae': ae_metric.compute().item(),
+        'acc': acc_metric.compute().item(),
+        'f1': f1_metric.compute().item(),
+        'prec': prec_metric.compute().item(),
+        'rec': rec_metric.compute().item(),
+        'spec': spec_metric.compute().item(),
+        'dice': dice_metric.compute().item(),
+        'iou': iou_metric.compute().item(),
+        'ck': ck_metric.compute().item(),
+        'roc_auc': auroc_metric.compute().item(),
+        'loss': running_loss / len(eval_loader.sampler),
+        'frac_sim': avg_fracture_sim,
+    }
+
+
+def eval_single(gt, pred, threshold=0.5, device="cpu", ignore_index=None):
+    """Evaluate metrics for a single prediction and ground truth pair."""
+    gt = torch.from_numpy(gt).to(device).float().unsqueeze(0).unsqueeze(0)
+    pred = torch.from_numpy(pred).to(device).float().unsqueeze(0).unsqueeze(0)
+
+    pred_clf = (pred > threshold).long()
+    gt_clf = (gt > 0).long()
+    if ignore_index not in [0, 1]:
+        ignore_index = None
+
+    # Metrics
+    acc_metric = BinaryAccuracy(ignore_index=ignore_index).to(device)
+    f1_metric = BinaryF1Score(ignore_index=ignore_index).to(device)
+    prec_metric = BinaryPrecision(ignore_index=ignore_index).to(device)
+    rec_metric = BinaryRecall(ignore_index=ignore_index).to(device)
+    spec_metric = BinarySpecificity(ignore_index=ignore_index).to(device)
+    auroc_metric = BinaryAUROC(ignore_index=ignore_index).to(device)
+    iou_metric = BinaryJaccardIndex(ignore_index=ignore_index).to(device)
+    dice_metric = DiceScore(num_classes=1, average="micro").to(device)
+    ck_metric = BinaryCohenKappa().to(device)
+    mse_metric = MeanSquaredError().to(device)
+    ae_metric = MeanAbsoluteError().to(device)
+    psnr_metric = PeakSignalNoiseRatio(data_range=1.0).to(device)
+    ssim_metric = StructuralSimilarityIndexMeasure().to(device)
+
+    # Compute metrics
+    acc_metric(pred_clf, gt_clf)
+    f1_metric(pred_clf, gt_clf)
+    prec_metric(pred_clf, gt_clf)
+    rec_metric(pred_clf, gt_clf)
+    spec_metric(pred_clf, gt_clf)
+    if gt_clf.numel() > 0 and gt_clf.min() != gt_clf.max():
+        auroc_metric(pred, gt_clf.int())
+    dice_metric(pred_clf, gt_clf)
+    iou_metric(pred_clf, gt_clf)
+    ck_metric(pred_clf, gt_clf)
+    mse_metric(pred, gt)
+    psnr_metric(pred, gt)
+    ssim_metric(pred, gt)
+    ae_metric(pred, gt)
+
+    return {
+        'mse': mse_metric.compute().item(),
+        'psnr': psnr_metric.compute().item(),
+        'ssim': ssim_metric.compute().item(),
+        'ae': ae_metric.compute().item(),
+        'acc': acc_metric.compute().item(),
+        'f1': f1_metric.compute().item(),
+        'prec': prec_metric.compute().item(),
+        'rec': rec_metric.compute().item(),
+        'spec': spec_metric.compute().item(),
+        'dice': dice_metric.compute().item(),
+        'iou': iou_metric.compute().item(),
+        'ck': ck_metric.compute().item(),
+        'roc_auc': auroc_metric.compute().item(),
+    }
+
+
+def save_metrics(metrics: dict, kind: str, writer, epoch: int):
+    """Log metrics to a TensorBoard writer."""
+    writer.add_scalar(f"Loss/{kind}", metrics['loss'], epoch)
+    writer.add_scalar(f"ACC/{kind}", metrics['acc'], epoch)
+    writer.add_scalar(f"F1/{kind}", metrics['f1'], epoch)
+    writer.add_scalar(f"PREC/{kind}", metrics['prec'], epoch)
+    writer.add_scalar(f"REC/{kind}", metrics['rec'], epoch)
+    writer.add_scalar(f"ROC_AUC/{kind}", metrics['roc_auc'], epoch)
+    writer.add_scalar(f"MSE/{kind}", metrics['mse'], epoch)
+    writer.add_scalar(f"PSNR/{kind}", metrics['psnr'], epoch)
+    writer.add_scalar(f"SSIM/{kind}", metrics['ssim'], epoch)
+    writer.add_scalar(f"SPEC/{kind}", metrics['spec'], epoch)
+    writer.add_scalar(f"DICE/{kind}", metrics['dice'], epoch)
+    writer.add_scalar(f"AE/{kind}", metrics['ae'], epoch)
+    writer.add_scalar(f"IoU/{kind}", metrics['iou'], epoch)
+    writer.flush()