rank_metric_with_precision

import numpy as np
import pickle as pkl
import os
import cv2


def calc_iou(mask_a, mask_b):
    intersection = (mask_a + mask_b >= 2).astype(np.float32).sum()
    iou = intersection / (mask_a + mask_b >= 1).astype(np.float32).sum()
    return iou


def get_cmp_mat(matrix):
    mat1, mat2 = np.meshgrid(matrix, matrix)
    com_mat = mat2 - mat1
    com_mat = (np.triu(com_mat) > 0).astype(np.int)
    return com_mat


def match(matrix, iou_thread, img_name):
    matched_gts = np.arange(matrix.shape[0])
    matched_ranks = matrix.argsort()[:, -1]
    for i, j in zip(matched_gts, matched_ranks):
        if matrix[i][j] < iou_thread:
            matched_ranks[i] = -1
    if len(set(matched_ranks[matched_ranks != -1])) < len(matched_ranks[matched_ranks != -1]):
        print(img_name)
        raise KeyError
    return matched_ranks


def main(result_dir, iou_thread):
    with open(result_dir, 'rb') as f:
        results = pkl.load(f)
    precision_sum = 0
    num = len(results)
    for indx, result in enumerate(results):
        gt_masks = result['gt_masks']
        segmaps = result['segmaps']
        gt_ranks = result['gt_ranks']
        rank_scores = result['rank_scores']
        name = result['img_name']

        if len(gt_ranks) == 1:
            num = num - 1
            continue
        gt_com_mat = get_cmp_mat(gt_ranks)
        if len(segmaps) == 0:
            rank_scores = np.zeros_like(gt_ranks)
        else:
            segmaps[segmaps > 0.5] = 1
            segmaps[segmaps <= 0.5] = 0

            ious = np.zeros([len(gt_masks), len(segmaps)])
            for i in range(len(gt_masks)):
                for j in range(len(segmaps)):
                    ious[i][j] = calc_iou(gt_masks[i], segmaps[j])
            matched_ranks = match(ious, iou_thread, name)
            unmatched_index = np.argwhere(matched_ranks == -1).squeeze(1)
            matched_ranks = matched_ranks[matched_ranks != -1]
            rank_scores = rank_scores[matched_ranks]
            for i in range(len(unmatched_index)):
                rank_scores = np.insert(rank_scores, unmatched_index[i], 0)
        rank_com_mat = get_cmp_mat(rank_scores)
        comp = np.abs(rank_com_mat - gt_com_mat)
        n = len(gt_ranks)
        precision = 1 - np.sum(comp)/(n*(n-1)/2)
        precision_sum += precision
        print(indx, name, precision)

    final_precision = precision_sum / num
    print(final_precision)


if __name__ == '__main__':
    root_dir = r'D:\dataset\salient_instance\SOC6K\soc_rank_saliency\test'
    result = '4_80000'
    result_dir = os.path.join(root_dir, '{}.pkl'.format(result))
    iou_thread = 0.5
    main(result_dir, iou_thread)