| """ |
| Copyright © 2025 Howard Hughes Medical Institute, Authored by Carsen Stringer , Michael Rariden and Marius Pachitariu. |
| """ |
| import numpy as np |
| from scipy.optimize import linear_sum_assignment |
| from scipy.sparse import csr_matrix |
|
|
|
|
| def average_precision(masks_true, masks_pred, threshold=[0.5, 0.75, 0.9]): |
| """ |
| Average precision estimation: AP = TP / (TP + FP + FN) |
| |
| This function is based heavily on the *fast* stardist matching functions |
| (https://github.com/mpicbg-csbd/stardist/blob/master/stardist/matching.py) |
| |
| Args: |
| masks_true (list of np.ndarrays (int) or np.ndarray (int)): |
| where 0=NO masks; 1,2... are mask labels |
| masks_pred (list of np.ndarrays (int) or np.ndarray (int)): |
| np.ndarray (int) where 0=NO masks; 1,2... are mask labels |
| |
| Returns: |
| ap (array [len(masks_true) x len(threshold)]): |
| average precision at thresholds |
| tp (array [len(masks_true) x len(threshold)]): |
| number of true positives at thresholds |
| fp (array [len(masks_true) x len(threshold)]): |
| number of false positives at thresholds |
| fn (array [len(masks_true) x len(threshold)]): |
| number of false negatives at thresholds |
| """ |
| not_list = False |
| if not isinstance(masks_true, list): |
| masks_true = [masks_true] |
| masks_pred = [masks_pred] |
| not_list = True |
| if not isinstance(threshold, list) and not isinstance(threshold, np.ndarray): |
| threshold = [threshold] |
|
|
| if len(masks_true) != len(masks_pred): |
| raise ValueError( |
| "metrics.average_precision requires len(masks_true)==len(masks_pred)") |
|
|
| ap = np.zeros((len(masks_true), len(threshold)), np.float32) |
| tp = np.zeros((len(masks_true), len(threshold)), np.float32) |
| fp = np.zeros((len(masks_true), len(threshold)), np.float32) |
| fn = np.zeros((len(masks_true), len(threshold)), np.float32) |
| n_true = np.array([len(np.unique(mt)) - 1 for mt in masks_true]) |
| n_pred = np.array([len(np.unique(mp)) - 1 for mp in masks_pred]) |
|
|
| for n in range(len(masks_true)): |
| |
| if n_pred[n] > 0: |
| iou = _intersection_over_union(masks_true[n], masks_pred[n])[1:, 1:] |
| for k, th in enumerate(threshold): |
| tp[n, k] = _true_positive(iou, th) |
| fp[n] = n_pred[n] - tp[n] |
| fn[n] = n_true[n] - tp[n] |
| ap[n] = tp[n] / (tp[n] + fp[n] + fn[n]) |
|
|
| if not_list: |
| ap, tp, fp, fn = ap[0], tp[0], fp[0], fn[0] |
| return ap, tp, fp, fn |
|
|
|
|
| def _intersection_over_union(masks_true, masks_pred): |
| """Calculate the intersection over union of all mask pairs. |
| |
| Parameters: |
| masks_true (np.ndarray, int): Ground truth masks, where 0=NO masks; 1,2... are mask labels. |
| masks_pred (np.ndarray, int): Predicted masks, where 0=NO masks; 1,2... are mask labels. |
| |
| Returns: |
| iou (np.ndarray, float): Matrix of IOU pairs of size [x.max()+1, y.max()+1]. |
| |
| How it works: |
| The overlap matrix is a lookup table of the area of intersection |
| between each set of labels (true and predicted). The true labels |
| are taken to be along axis 0, and the predicted labels are taken |
| to be along axis 1. The sum of the overlaps along axis 0 is thus |
| an array giving the total overlap of the true labels with each of |
| the predicted labels, and likewise the sum over axis 1 is the |
| total overlap of the predicted labels with each of the true labels. |
| Because the label 0 (background) is included, this sum is guaranteed |
| to reconstruct the total area of each label. Adding this row and |
| column vectors gives a 2D array with the areas of every label pair |
| added together. This is equivalent to the union of the label areas |
| except for the duplicated overlap area, so the overlap matrix is |
| subtracted to find the union matrix. |
| """ |
| if masks_true.size != masks_pred.size: |
| raise ValueError(f"masks_true.size {masks_true.shape} != masks_pred.size {masks_pred.shape}") |
| overlap = csr_matrix((np.ones((masks_true.size,), "int"), |
| (masks_true.flatten(), masks_pred.flatten())), |
| shape=(masks_true.max()+1, masks_pred.max()+1)) |
| overlap = overlap.toarray() |
| n_pixels_pred = np.sum(overlap, axis=0, keepdims=True) |
| n_pixels_true = np.sum(overlap, axis=1, keepdims=True) |
| iou = overlap / (n_pixels_pred + n_pixels_true - overlap) |
| iou[np.isnan(iou)] = 0.0 |
| return iou |
|
|
|
|
| def _true_positive(iou, th): |
| """Calculate the true positive at threshold th. |
| |
| Args: |
| iou (float, np.ndarray): Array of IOU pairs. |
| th (float): Threshold on IOU for positive label. |
| |
| Returns: |
| tp (float): Number of true positives at threshold. |
| |
| How it works: |
| (1) Find minimum number of masks. |
| (2) Define cost matrix; for a given threshold, each element is negative |
| the higher the IoU is (perfect IoU is 1, worst is 0). The second term |
| gets more negative with higher IoU, but less negative with greater |
| n_min (but that's a constant...). |
| (3) Solve the linear sum assignment problem. The costs array defines the cost |
| of matching a true label with a predicted label, so the problem is to |
| find the set of pairings that minimizes this cost. The scipy.optimize |
| function gives the ordered lists of corresponding true and predicted labels. |
| (4) Extract the IoUs from these pairings and then threshold to get a boolean array |
| whose sum is the number of true positives that is returned. |
| """ |
| n_min = min(iou.shape[0], iou.shape[1]) |
| costs = -(iou >= th).astype(float) - iou / (2 * n_min) |
| true_ind, pred_ind = linear_sum_assignment(costs) |
| match_ok = iou[true_ind, pred_ind] >= th |
| tp = match_ok.sum() |
| return tp |
|
|
