compute_metric.py

"""
Created on Thu Mar 31 18:10:52 2022
adapted form https://github.com/stardist/stardist/blob/master/stardist/matching.py
Thanks the authors of Stardist for sharing the great code

"""

import argparse
from multiprocessing.sharedctypes import Value
from typing import Tuple
import numpy as np
from numba import jit
from scipy.optimize import linear_sum_assignment
from collections import OrderedDict
import pandas as pd
from skimage import segmentation
import tifffile as tif
import os
join = os.path.join
from tqdm import tqdm
import torch
from torch.nn import functional as F
from itertools import product

def _intersection_over_union(masks_true, masks_pred):
    """ intersection over union of all mask pairs
    
    Parameters
    ------------
    
    masks_true: ND-array, int 
        ground truth masks, where 0=NO masks; 1,2... are mask labels
    masks_pred: ND-array, int
        predicted masks, where 0=NO masks; 1,2... are mask labels
    """
    overlap = _label_overlap(masks_true, masks_pred)
    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

@jit(nopython=True)
def _label_overlap(x, y):
    """ fast function to get pixel overlaps between masks in x and y 
    
    Parameters
    ------------

    x: ND-array, int
        where 0=NO masks; 1,2... are mask labels
    y: ND-array, int
        where 0=NO masks; 1,2... are mask labels

    Returns
    ------------

    overlap: ND-array, int
        matrix of pixel overlaps of size [x.max()+1, y.max()+1]
    
    """
    x = x.ravel()
    y = y.ravel()
    
    # preallocate a 'contact map' matrix
    overlap = np.zeros((1+x.max(),1+y.max()), dtype=np.uint)
    
    # loop over the labels in x and add to the corresponding
    # overlap entry. If label A in x and label B in y share P
    # pixels, then the resulting overlap is P
    # len(x)=len(y), the number of pixels in the whole image 
    for i in range(len(x)):
        overlap[x[i],y[i]] += 1
    return overlap

def _true_positive(iou, th):
    """ true positive at threshold th
    
    Parameters
    ------------

    iou: float, ND-array
        array of IOU pairs
    th: float
        threshold on IOU for positive label

    Returns
    ------------

    tp: float
        number of true positives at threshold
    """
    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

def eval_tp_fp_fn(masks_true, masks_pred, threshold=0.5):
    num_inst_gt = np.max(masks_true)
    num_inst_seg = np.max(masks_pred)
    if num_inst_seg>0:
        iou = _intersection_over_union(masks_true, masks_pred)[1:, 1:]
            # for k,th in enumerate(threshold):
        tp = _true_positive(iou, threshold)
        fp = num_inst_seg - tp
        fn = num_inst_gt - tp
    else:
        print('No segmentation results!')
        tp = 0
        fp = 0
        fn = 0
        
    return tp, fp, fn

def remove_boundary_cells(mask):
    W, H = mask.shape
    bd = np.ones((W, H))
    bd[2:W-2, 2:H-2] = 0
    bd_cells = np.unique(mask*bd)
    for i in bd_cells[1:]:
        mask[mask==i] = 0
    new_label,_,_ = segmentation.relabel_sequential(mask)
    return new_label

def eval_one(gt_path, seg_path, name):
    # Load the images for this case
    gt = tif.imread(join(gt_path, name))
    seg = tif.imread(join(seg_path, name))
    # Score the cases
    # do not consider cells on the boundaries during evaluation
    if np.prod(gt.shape)<25000000:
        gt = remove_boundary_cells(gt.astype(np.int32)) 
        seg = remove_boundary_cells(seg.astype(np.int32))           
        tp, fp, fn = eval_tp_fp_fn(gt, seg, threshold=0.5)
    else: # for large images (>5000x5000), the F1 score is computed by a patch-based way
        H, W = gt.shape
        roi_size = 2000
    
        if H % roi_size != 0:
            n_H = H // roi_size + 1
            new_H = roi_size * n_H
        else:
            n_H = H // roi_size
            new_H = H
    
        if W % roi_size != 0:
            n_W = W // roi_size + 1
            new_W = roi_size * n_W    
        else:
            n_W = W // roi_size
            new_W = W    
    
        gt_pad = np.zeros((new_H, new_W), dtype=gt.dtype)
        seg_pad = np.zeros((new_H, new_W), dtype=gt.dtype)
        gt_pad[:H, :W] = gt
        seg_pad[:H, :W] = seg
            
        tp = 0
        fp = 0
        fn = 0
        for i in range(n_H):
            for j in range(n_W):
                gt_roi  = remove_boundary_cells(gt_pad[roi_size*i:roi_size*(i+1), roi_size*j:roi_size*(j+1)])
                seg_roi = remove_boundary_cells(seg_pad[roi_size*i:roi_size*(i+1), roi_size*j:roi_size*(j+1)])
                tp_i, fp_i, fn_i = eval_tp_fp_fn(gt_roi, seg_roi, threshold=0.5)
                tp += tp_i
                fp += fp_i
                fn += fn_i            
    
    if tp == 0:
        precision = 0
        recall = 0
        f1 = 0
    else:
        precision = tp / (tp + fp)
        recall = tp / (tp + fn)
        f1 = 2*(precision * recall)/ (precision + recall)

    return name, np.round(f1, 4)


def generate_match(
    masks_true: np.ndarray,
    masks_pred: np.ndarray,
    single_th: float=0.5
):
    num_inst_seg = np.max(masks_pred)
    if num_inst_seg == 0:
        return np.zeros([0, 2]) # K x 2 array
    iou = _intersection_over_union(masks_true, masks_pred)[1:, 1:]

    n_min = min(iou.shape[0], iou.shape[1])
    costs = -(iou >= single_th).astype(float) - iou / (2*n_min)
    true_ind, pred_ind = linear_sum_assignment(costs)
    match_ok = iou[true_ind, pred_ind] >= single_th

    true_ind = true_ind[match_ok]
    pred_ind = pred_ind[match_ok]

    return np.stack([true_ind, pred_ind], axis=1) + 1 # K x 2 array

def eval_tp_fp_fn_multi(
    masks_true: np.ndarray, 
    masks_pred: np.ndarray, 
    threshold: np.ndarray) -> Tuple[np.ndarray]:

    num_inst_gt = np.max(masks_true)
    num_inst_seg = np.max(masks_pred)
    if num_inst_seg == 0:
        return tuple(np.zeros((3, threshold.shape[0]), dtype=np.int32))

    iou = _intersection_over_union(masks_true, masks_pred)[1:, 1:]
    res = []
    for single_th in threshold:
        n_min = min(iou.shape[0], iou.shape[1])
        costs = -(iou >= single_th).astype(float) - iou / (2*n_min)
        true_ind, pred_ind = linear_sum_assignment(costs)
        match_ok = iou[true_ind, pred_ind] >= single_th
        res.append(match_ok.sum())
    tp = np.array(res)
    fp = num_inst_seg - tp
    fn = num_inst_gt - tp
    return tp, fp, fn

def eval_one_multithres(
    pred: np.ndarray, 
    gt: np.ndarray,
    *,
    th = None):
    if th is None:
        th = np.linspace(0.5, 0.95, 10) # follow mAP0.5:0.95
    if gt.shape != pred.shape or len(gt.shape) != 2:
        raise ValueError(f'Incorperate mask shape with prediction `{pred.shape}`, while mask `{gt.shape}`')
    n_H, n_W, H, W = 1, 1, *gt.shape
    roi_size = 2000
    if np.prod(gt.shape) >= 25000000:
        if H % roi_size != 0:
            n_H = H // roi_size + 1
            new_H = roi_size * n_H
        else:
            n_H = H // roi_size
            new_H = H
    
        if W % roi_size != 0:
            n_W = W // roi_size + 1
            new_W = roi_size * n_W    
        else:
            n_W = W // roi_size
            new_W = W    
    
        gt = np.pad(gt, ((0, new_H - H), (0, new_W - W)))
        pred = np.pad(pred, ((0, new_H - H), (0, new_W - W)))
            
    tp, fp, fn = 0, 0, 0
    if max(n_H, n_W) > 1:
        for i, j in product(range(n_H), range(n_W)):
                gt_roi  = remove_boundary_cells(gt[roi_size*i:roi_size*(i+1), roi_size*j:roi_size*(j+1)])
                seg_roi = remove_boundary_cells(pred[roi_size*i:roi_size*(i+1), roi_size*j:roi_size*(j+1)])
                tp_i, fp_i, fn_i = eval_tp_fp_fn_multi(gt_roi, seg_roi, th)
                tp, fp, fn = tp + tp_i, fp + fp_i, fn + fn_i
    else:
        gt_roi, seg_roi  = remove_boundary_cells(gt), remove_boundary_cells(pred)
        tp, fp, fn = eval_tp_fp_fn_multi(gt_roi, seg_roi, th)


    precision = tp / (tp + fp + 1e-10)
    recall = tp / (tp + fn + 1e-10)
    f1 = 2*(precision * recall)/ (precision + recall + 1e-10)
    return np.stack([precision, recall, f1, th], axis=-1)

def main():
    from utils.taskrunner import TaskRunner
    parser = argparse.ArgumentParser('Compute F1 score for cell segmentation results', add_help=False)
    # Dataset parameters
    parser.add_argument('--gt_path', type=str, help='path to ground truth; file names end with _label.tiff', required=True)
    parser.add_argument('--seg_path', type=str, help='path to segmentation results; file names are the same as ground truth', required=True)
    parser.add_argument('--save_path', default=None, help='path where to save metrics')
    args = parser.parse_args()

    gt_path = args.gt_path
    seg_path = args.seg_path
    names = sorted(os.listdir(seg_path))
    seg_metric = OrderedDict()
    seg_metric['Names'] = []
    seg_metric['F1_Score'] = []

    runner = TaskRunner(12)

    all_args = []
    for name in names:
        if not name.endswith('_label.tiff'):
            continue
        all_args.append({
            'gt_path': gt_path,
            'seg_path': seg_path,
            'name': name
        })

    results = runner.run(eval_one, all_args, total=len(all_args))
    for name, f1 in results:
        seg_metric['Names'].append(name)
        seg_metric['F1_Score'].append(f1)
    
    if args.save_path is not None:
        seg_metric_df = pd.DataFrame(seg_metric)
        seg_metric_df.to_csv(join(args.save_path, 'seg_metric.csv'), index=False)
    print('mean F1 Score:', np.mean(seg_metric['F1_Score']))
    

if __name__ == '__main__':
    main()