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)