app.py

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
from PIL import Image
from segmentation_models_pytorch import Segformer
from skimage import color, io
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
# ------------------------------------------------------------
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")


# ------------------------------------------------------------
# Canny edge detection
# ------------------------------------------------------------
def canny_fn(img, sigma, lt, ht):
    """Apply Canny edge detection using skimage."""
    if img is None:
        return None

    if img.ndim == 3:
        gray = color.rgb2gray(img)
    else:
        gray = img.astype(float) / 255.

    edges = canny(
        gray,
        sigma=sigma,
        low_threshold=lt,
        high_threshold=ht
    )
    return (255 - edges * 255).astype(np.uint8)


# ------------------------------------------------------------
# Phase Congruency (phasepack)
# ------------------------------------------------------------
def phase_congruency_fn(
    x,
    img,
    nscale,
    norient,
    minWaveLength,
    mult,
    sigmaOnf,
    k,
    cutOff,
    g,
    noiseMethod,
):
    """Compute phase congruency with adjustable parameters."""
    if img is None:
        return None

    if img.ndim == 3:
        gray = color.rgb2gray(img)
    else:
        gray = img.astype(float) / 255.

    pc, m, ori, ft, PC, EO, T = phasecong(
        gray,
        nscale=nscale,
        norient=norient,
        minWaveLength=minWaveLength,
        mult=mult,
        sigmaOnf=sigmaOnf,
        k=k,
        cutOff=cutOff,
        g=g,
        noiseMethod=noiseMethod,
    )

    # Threshold using slider
    pc = pc < x
    return (pc * 255).astype(np.uint8)


# ------------------------------------------------------------
# Sato vesselness-like filter
# ------------------------------------------------------------
sato_sigmas_list = [
    range(1, 5),
    range(1, 20, 4),
    (2,),
    (1,),
]


def sato_fn(img, x, sigmas):
    """Sato ridge detection over selected sigma set."""
    if img is None:
        return None
    gray = color.rgb2gray(img)
    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=int(threshold*255),
                          device=device)

    df = pd.DataFrame([metrics])
    df = df.round(3)
    return df


# ------------------------------------------------------------
# Deep learning model loading
# ------------------------------------------------------------
def load_model(model_name: str):
    """Load segmentation model weights."""
    if model_name.lower() == "unet":
        model = UNet(init_features=64)
        weight_path = "model/unet.pt"

    elif model_name.lower() == "segformer":
        model = Segformer(
            encoder_name='resnet34',
            encoder_depth=5,
            encoder_weights='imagenet',
            decoder_segmentation_channels=256,
            in_channels=4,
            classes=1,
            activation='sigmoid'
        )
        weight_path = "model/segformer.pt"

    else:
        raise ValueError(f"Unknown model: {model_name}")

    model.load_state_dict(
        torch.load(weight_path, weights_only=True, map_location=torch.device('cpu'))
    )
    model.to(device)
    model.eval()

    return model


# ------------------------------------------------------------
# Inference on RGB + DEM pair
# ------------------------------------------------------------
def run_inference(img_path, dem_path, model_name):
    """Run patch-based inference for fracture segmentation."""
    model = load_model(model_name)

    img = io.imread(img_path)
    dem = io.imread(dem_path)

    # Ensure RGB format
    if img.ndim == 2:
        img = np.stack([img, img, img], axis=-1)
    if img.shape[2] > 3:
        img = img[:, :, :3]

    # Merge RGB + DEM
    combined = np.concatenate((img[:, :, :3], np.expand_dims(dem, 2)), 2)

    patch_shape = 256
    h, w, c = combined.shape

    # Padding for patchify
    pad_h = (patch_shape - h % patch_shape) % patch_shape
    pad_w = (patch_shape - w % patch_shape) % patch_shape

    combined_padded = np.pad(
        combined,
        ((0, pad_h), (0, pad_w), (0, 0)),
        mode="constant",
        constant_values=0,
    )

    # Patchify
    patches = patchify(
        combined_padded,
        (patch_shape, patch_shape, c),
        step=patch_shape,
    )

    pred_patches = []
    for i in range(patches.shape[0]):
        for j in range(patches.shape[1]):
            single_patch = patches[i, j, :, :, :, :]
            single_patch = torch.Tensor(np.array(single_patch))
            single_patch = single_patch.permute(0, 3, 1, 2) / 255.

            with torch.no_grad():
                patch_pred = model(single_patch.to(device))

            pred_patches.append(patch_pred.cpu())

    # Reshape back to full image
    pred = np.array(pred_patches)
    pred = np.reshape(
        pred,
        (patches.shape[0], patches.shape[1], 1, patch_shape, patch_shape, 1),
    )
    pred = unpatchify(pred, combined_padded.shape[:2] + (1,))

    pred = pred[:h, :w, :]
    pred = (255 - pred * 255).astype(np.uint8)

    return Image.fromarray(img[:, :, :3]), Image.fromarray(pred.reshape(h, w))


# ------------------------------------------------------------
# User Interface
# ------------------------------------------------------------
with gr.Blocks(title="Fractex2D Segmentation") as demo:

    gr.Markdown("# **Fractex2D – Fracture Detection**")

    gr.Markdown(
        """
        Try out deep models that use RGB+DEM inputs along with classic vision methods that work on RGB images.
        Support for RGB-only deep models is on the way.
        """
    )

    with gr.Row():

        # ------------------------------------------------------------
        # TAB 1 — DEEP LEARNING
        # ------------------------------------------------------------
        with gr.Tab("DEEP LEARNING"):
            gr.Markdown(
                """
                ## Deep Learning Segmentation  
                Patch-based fracture segmentation using **UNet** or **SegFormer** trained on [FraXet]() dataset.

                **Requirements before running:**
                - RGB image: `.png` or `.tif`
                - DEM: `.tif`
                - Both must have **same resolution**

                The model processes the RGB + DEM pair in 256×256 patches internally to produce a  binary fracture map, while still allowing you to **input images of any size**.
                """
            )

            with gr.Row():
                rgb_input = gr.File(type="filepath", label="RGB image (.png/.tif)")
                dem_input = gr.File(type="filepath", label="DEM (.tif)")
                model_choice = gr.Dropdown(
                    choices=["unet", "segformer"],
                    value="segformer",
                    label="Model",
                )

            with gr.Row():
                with gr.Column(scale=1):  # empty column to push btn to center
                    pass
                run_btn = gr.Button("Run Segmentation", elem_id="run-button")
                with gr.Column(scale=1):  # empty column to balance
                    pass

            with gr.Row():
                rgb_show = gr.Image(type="pil", label="Input RGB")
                pred_show = gr.Image(type="pil", label="Prediction")

            gr.Examples(
                examples=[
                    ["examples/kl5-s3_1.png", "examples/kl5-s3-dem_1.tif", "unet"],
                    ["examples/kl5-s3_1.png", "examples/kl5-s3-dem_1.tif", "segformer"],
                ],
                inputs=[rgb_input, dem_input, model_choice],
            )

            run_btn.click(
                fn=run_inference,
                inputs=[rgb_input, dem_input, model_choice],
                outputs=[rgb_show, pred_show],
            )

        # ------------------------------------------------------------
        # TAB 2 — SATO FILTER
        # ------------------------------------------------------------
        with gr.Tab("Sato"):
            gr.Markdown(
                """
                ## Sato Ridge Detection  
                Vesselness-inspired filter (scikit-image) useful for enhancing elongated structures https://doi.org/10.1016/S1361-8415(98)80009-1.
                Adjust threshold and σ-sets to explore different ridge responses.
                """
            )

            with gr.Row():
                with gr.Column(scale=1):
                    sato_in = gr.Image(value="examples/kl5-s3_1.png")
                    sato_x = gr.Slider(0, 1, value=0.08, step=0.01, label="Threshold")
                    sato_sigmas = gr.Radio(
                        [('range(1,5)', 0),
                         ('range(1,20,4)', 1),
                         ('(2,)', 2),
                         ('(1,)', 3)],
                        label="Sigma set",
                        value=0,
                    )
                with gr.Column(scale=1):
                    sato_out = gr.Image()

        # Auto update
        sato_in.change(sato_fn, [sato_in, sato_x, sato_sigmas], sato_out)
        sato_x.change(sato_fn, [sato_in, sato_x, sato_sigmas], sato_out)
        sato_sigmas.change(sato_fn, [sato_in, sato_x, sato_sigmas], sato_out)

        # ------------------------------------------------------------
        # TAB 3 — CANNY
        # ------------------------------------------------------------
        with gr.Tab("Canny"):
            gr.Markdown(
                """
                ## Canny edge detection
                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
                """
            )

            with gr.Row():
                with gr.Column(scale=1):
                    canny_in = gr.Image(value="examples/kl5-s3_1.png")

                    canny_sigma = gr.Slider(
                        0, 7,
                        value=1.37,
                        step=0.01,
                        label="Sigma"
                    )

                    canny_lt = gr.Slider(
                        0, 1,
                        value=0.37,
                        step=0.01,
                        label="Low threshold"
                    )

                    canny_ht = gr.Slider(
                        0, 1,
                        value=0.58,
                        step=0.01,
                        label="High threshold"
                    )

                with gr.Column(scale=1):
                    canny_out = gr.Image()

        canny_in.change(canny_fn, [canny_in, canny_sigma, canny_lt, canny_ht], canny_out)
        canny_sigma.change(canny_fn, [canny_in, canny_sigma, canny_lt, canny_ht], canny_out)
        canny_lt.change(canny_fn, [canny_in, canny_sigma, canny_lt, canny_ht], canny_out)
        canny_ht.change(canny_fn, [canny_in, canny_sigma, canny_lt, canny_ht], canny_out)

        # ------------------------------------------------------------
        # TAB 4 — PHASE CONGRUENCY
        # ------------------------------------------------------------
        with gr.Tab("Phase Congruency"):
            gr.Markdown(
                """
                ## Phase Congruency
                Edge/line detection ([phasepack](https://github.com/alimuldal/phasepack)) based on phase agreement in the frequency domain https://doi.org/10.1007/s004260000024.  
                Computationally expensive → runs **only on button click**.
                Useful for illumination-invariant structural extraction.
                """
            )

            with gr.Row():
                pc_in = gr.Image(value="examples/kl5-s3_1.png")
                pc_out = gr.Image()

            with gr.Row():
                with gr.Column(scale=1):  # empty column to push btn to center
                    pass
                pc_btn = gr.Button("Detect")
                with gr.Column(scale=1):  # empty column to balance
                    pass

            with gr.Row():
                with gr.Column(scale=1):
                    x_pc = gr.Slider(0, 1, value=0.15, step=0.01, label="Threshold")
                    pc_nscale = gr.Slider(3, 10, value=6, step=1, label="nscale")
                    pc_norient = gr.Slider(1, 16, value=8, step=1, label="norient")
                    pc_minWL = gr.Slider(1, 10, value=4, step=1, label="minWaveLength")
                    pc_mult = gr.Slider(1.0, 5.0, value=2.1, step=0.1, label="mult")
                with gr.Column(scale=1):
                    pc_sigma = gr.Slider(0.1, 1.0, value=0.35, step=0.05, label="sigmaOnf")
                    pc_k = gr.Slider(0.1, 10.0, value=2.8, step=0.1, label="k")
                    pc_cutoff = gr.Slider(0.0, 1.0, value=0.5, step=0.01, label="cutOff")
                    pc_g = gr.Slider(0.1, 50.0, value=10.6, step=0.5, label="g")
                    pc_noise = gr.Slider(-2, 2, value=-1, step=1, label="noiseMethod")

            pc_btn.click(
                fn=phase_congruency_fn,
                inputs=[
                    x_pc, pc_in, pc_nscale, pc_norient, pc_minWL, pc_mult,
                    pc_sigma, pc_k, pc_cutoff, pc_g, pc_noise
                ],
                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
    # ------------------------------------------------------------
    gr.Markdown(
        """
        The sample images included with this interface originate from: 
        Nordbäck, N., & Ovaskainen, N. (2022). UAV-acquired orthomosaics of \
        Loviisa shoreline outcrops (Version 1.0.0) [Dataset]. Zenodo. \
        https://doi.org/10.5281/zenodo.7077519
        """
    )

if __name__ == "__main__":
    demo.launch()