AI Workflows

## Task provided by user

Create a Python script that calculates 3D segmentation accuracy. The inputs are two 3D binary layers, and the accuracy should be calculated as follows:

Object Segmentation Accuracy

Calculates the average precision to evaluate AI performance on objects. This node has two inputs - GT (Ground Truth) and Pred (Prediction). It compares the ground truth binary layer (A) and predicted binary layer (B) generated by segmentation using AI.
This node calculates the given metrics directly on objects. The metrics it computes (TP, FP, FN) are expressed as the number of objects.
- true positives (TP) matched correctly,
- false positives (FP) incorrectly segmented objects and
- false negatives (FN) incorrectly missed objects.
Based on these numbers it calculates:
- precision = TP / (TP + FP),
- recall = TP / (TP + FN) and
- F1 = 2 x precision x recall / (precision + recall)

I need the accuracy to be calculated on complete 3D objects, not separately for individual Z-slices.

import limnode
import numpy as np
import pandas as pd
from scipy import ndimage
from scipy.sparse import csr_matrix
from scipy.sparse.csgraph import maximum_bipartite_matching


"""
3D object segmentation accuracy.

inp[0] = Ground Truth Binary
inp[1] = Prediction Binary
out[0] = Table

The complete input volume (Z, Y, X, C) is processed as one 3D volume.
Connected components therefore extend across Z slices and are NOT evaluated
slice-by-slice.

Objects are matched one-to-one when their 3D IoU >= IOU_THRESHOLD.
"""

IOU_THRESHOLD = 0.5


def output(
    inp: limnode.InDefTuple,
    out: limnode.OutDefTuple,
    par: limnode.UserParTuple,
) -> None:
    if len(inp) < 2:
        raise ValueError("Two Binary inputs are required: GT and Pred.")

    if not isinstance(inp[0], limnode.InputBinaryDef):
        raise TypeError("inp[0] must be a Binary input (Ground Truth).")

    if not isinstance(inp[1], limnode.InputBinaryDef):
        raise TypeError("inp[1] must be a Binary input (Prediction).")

    if len(out) < 1 or not isinstance(out[0], limnode.OutputTableDef):
        raise TypeError("out[0] must be a Table output.")

    (
        out[0]
        .makeEmpty("3D Object Segmentation Accuracy")
        .withLoopCols()
        .withIntCol("TP")
        .withIntCol("FP")
        .withIntCol("FN")
        .withFloatCol("Precision")
        .withFloatCol("Recall")
        .withFloatCol("F1")
    )


def build(
    par: limnode.UserParTuple,
    loops: limnode.LoopDefs,
) -> limnode.Program | None:
    # No reduction/multi-pass program is needed.
    # Each run() invocation receives the volume as (Z, Y, X, C).
    return None


def run(
    inp: limnode.InDataTuple,
    out: limnode.OutDataTuple,
    par: limnode.UserParTuple,
    ctx: limnode.RunContext,
) -> None:
    if len(inp) < 2:
        raise ValueError("Two Binary inputs are required: GT and Pred.")

    if not isinstance(inp[0], limnode.InputBinaryData):
        raise TypeError("inp[0] must be Binary data (Ground Truth).")

    if not isinstance(inp[1], limnode.InputBinaryData):
        raise TypeError("inp[1] must be Binary data (Prediction).")

    if len(out) < 1 or not isinstance(out[0], limnode.OutputTableData):
        raise TypeError("out[0] must be Table data.")

    if ctx.shouldAbort:
        return

    gt_data = np.asarray(inp[0].data)
    pred_data = np.asarray(inp[1].data)

    if gt_data.shape != pred_data.shape:
        raise ValueError(
            "GT and Pred must have identical shapes; "
            f"got {gt_data.shape} and {pred_data.shape}."
        )

    if gt_data.ndim != 4:
        raise ValueError(
            "Binary inputs must have shape (Z, Y, X, C)."
        )

    # Convert arbitrary binary bit depth to Boolean foreground.
    gt_mask = np.any(gt_data != 0, axis=3)
    pred_mask = np.any(pred_data != 0, axis=3)

    # 26-connectivity in a true 3D volume.
    structure = ndimage.generate_binary_structure(3, 3)

    gt_labels, n_gt = ndimage.label(
        gt_mask,
        structure=structure,
    )

    if ctx.shouldAbort:
        return

    pred_labels, n_pred = ndimage.label(
        pred_mask,
        structure=structure,
    )

    if ctx.shouldAbort:
        return

    tp = 0

    if n_gt > 0 and n_pred > 0:
        gt_sizes = np.bincount(
            gt_labels.ravel(),
            minlength=n_gt + 1,
        ).astype(np.int64, copy=False)

        pred_sizes = np.bincount(
            pred_labels.ravel(),
            minlength=n_pred + 1,
        ).astype(np.int64, copy=False)

        overlap_mask = (gt_labels > 0) & (pred_labels > 0)

        if np.any(overlap_mask):
            gt_ids_overlap = gt_labels[overlap_mask].astype(
                np.int64,
                copy=False,
            )
            pred_ids_overlap = pred_labels[overlap_mask].astype(
                np.int64,
                copy=False,
            )

            base = np.int64(n_pred + 1)

            pair_codes = gt_ids_overlap * base + pred_ids_overlap

            unique_codes, intersections = np.unique(
                pair_codes,
                return_counts=True,
            )

            gt_ids = unique_codes // base
            pred_ids = unique_codes % base

            unions = (
                gt_sizes[gt_ids]
                + pred_sizes[pred_ids]
                - intersections
            )

            iou = (
                intersections.astype(np.float64)
                / unions.astype(np.float64)
            )

            valid = iou >= IOU_THRESHOLD

            if np.any(valid):
                # Convert natural labels 1..N to sparse-matrix indices 0..N-1.
                rows = gt_ids[valid] - 1
                cols = pred_ids[valid] - 1

                graph = csr_matrix(
                    (
                        np.ones(
                            np.count_nonzero(valid),
                            dtype=np.uint8,
                        ),
                        (rows, cols),
                    ),
                    shape=(n_gt, n_pred),
                )

                # Maximum one-to-one matching.
                matching = maximum_bipartite_matching(
                    graph,
                    perm_type="column",
                )

                tp = int(np.count_nonzero(matching >= 0))

    fp = int(n_pred - tp)
    fn = int(n_gt - tp)

    precision = (
        float(tp / (tp + fp))
        if (tp + fp) > 0
        else 0.0
    )

    recall = (
        float(tp / (tp + fn))
        if (tp + fn) > 0
        else 0.0
    )

    f1 = (
        float(
            2.0 * precision * recall
            / (precision + recall)
        )
        if (precision + recall) > 0.0
        else 0.0
    )

    result = pd.DataFrame(
        {
            "TP": [tp],
            "FP": [fp],
            "FN": [fn],
            "Precision": [precision],
            "Recall": [recall],
            "F1": [f1],
        }
    )

    out[0].withDataFrame(
        result,
        [
            "TP",
            "FP",
            "FN",
            "Precision",
            "Recall",
            "F1",
        ],
    )


if __name__ == '__main__':
    from limnode import print
    limnode.child_main(run, output, build)