DIS/IS-Net/hce_metric_main.py

## hce_metric.py
import numpy as np
from skimage import io
import matplotlib.pyplot as plt
import cv2 as cv
from skimage.morphology import skeletonize
from skimage.morphology import erosion, dilation, disk
from skimage.measure import label

import os
import sys
from tqdm import tqdm
from glob import glob
import pickle as pkl

def filter_bdy_cond(bdy_, mask, cond):

    cond = cv.dilate(cond.astype(np.uint8),disk(1))
    labels = label(mask) # find the connected regions
    lbls = np.unique(labels) # the indices of the connected regions
    indep = np.ones(lbls.shape[0]) # the label of each connected regions
    indep[0] = 0 # 0 indicate the background region

    boundaries = []
    h,w = cond.shape[0:2]
    ind_map = np.zeros((h,w))
    indep_cnt = 0

    for i in range(0,len(bdy_)):
        tmp_bdies = []
        tmp_bdy = []
        for j in range(0,bdy_[i].shape[0]):
            r, c = bdy_[i][j,0,1],bdy_[i][j,0,0]

            if(np.sum(cond[r,c])==0 or ind_map[r,c]!=0):
                if(len(tmp_bdy)>0):
                    tmp_bdies.append(tmp_bdy)
                    tmp_bdy = []
                continue
            tmp_bdy.append([c,r])
            ind_map[r,c] =  ind_map[r,c] + 1
            indep[labels[r,c]] = 0 # indicates part of the boundary of this region needs human correction
        if(len(tmp_bdy)>0):
            tmp_bdies.append(tmp_bdy)

        # check if the first and the last boundaries are connected
        # if yes, invert the first boundary and attach it after the last boundary
        if(len(tmp_bdies)>1):
            first_x, first_y = tmp_bdies[0][0]
            last_x, last_y = tmp_bdies[-1][-1]
            if((abs(first_x-last_x)==1 and first_y==last_y) or
               (first_x==last_x and abs(first_y-last_y)==1) or
               (abs(first_x-last_x)==1 and abs(first_y-last_y)==1)
              ):
                tmp_bdies[-1].extend(tmp_bdies[0][::-1])
                del tmp_bdies[0]

        for k in range(0,len(tmp_bdies)):
            tmp_bdies[k] =  np.array(tmp_bdies[k])[:,np.newaxis,:]
        if(len(tmp_bdies)>0):
            boundaries.extend(tmp_bdies)

    return boundaries, np.sum(indep)

# this function approximate each boundary by DP algorithm
# https://en.wikipedia.org/wiki/Ramer%E2%80%93Douglas%E2%80%93Peucker_algorithm
def approximate_RDP(boundaries,epsilon=1.0):

    boundaries_ = []
    boundaries_len_ = []
    pixel_cnt_ = 0

    # polygon approximate of each boundary
    for i in range(0,len(boundaries)):
        boundaries_.append(cv.approxPolyDP(boundaries[i],epsilon,False))

    # count the control points number of each boundary and the total control points number of all the boundaries
    for i in range(0,len(boundaries_)):
        boundaries_len_.append(len(boundaries_[i]))
        pixel_cnt_ = pixel_cnt_ + len(boundaries_[i])

    return boundaries_, boundaries_len_, pixel_cnt_


def relax_HCE(gt, rs, gt_ske, relax=5, epsilon=2.0):
    # print("max(gt_ske): ", np.amax(gt_ske))
    # gt_ske = gt_ske>128
    # print("max(gt_ske): ", np.amax(gt_ske))

    # Binarize gt
    if(len(gt.shape)>2):
        gt = gt[:,:,0]

    epsilon_gt = 128#(np.amin(gt)+np.amax(gt))/2.0
    gt = (gt>epsilon_gt).astype(np.uint8)

    # Binarize rs
    if(len(rs.shape)>2):
        rs = rs[:,:,0]
    epsilon_rs = 128#(np.amin(rs)+np.amax(rs))/2.0
    rs = (rs>epsilon_rs).astype(np.uint8)

    Union = np.logical_or(gt,rs)
    TP = np.logical_and(gt,rs)
    FP = rs - TP
    FN = gt - TP

    # relax the Union of gt and rs
    Union_erode = Union.copy()
    Union_erode = cv.erode(Union_erode.astype(np.uint8),disk(1),iterations=relax)

    # --- get the relaxed False Positive regions for computing the human efforts in correcting them ---
    FP_ = np.logical_and(FP,Union_erode) # get the relaxed FP
    for i in range(0,relax):
        FP_ = cv.dilate(FP_.astype(np.uint8),disk(1))
        FP_ = np.logical_and(FP_, 1-np.logical_or(TP,FN))
    FP_ = np.logical_and(FP, FP_)

    # --- get the relaxed False Negative regions for computing the human efforts in correcting them ---
    FN_ = np.logical_and(FN,Union_erode) # preserve the structural components of FN
    ## recover the FN, where pixels are not close to the TP borders
    for i in range(0,relax):
        FN_ = cv.dilate(FN_.astype(np.uint8),disk(1))
        FN_ = np.logical_and(FN_,1-np.logical_or(TP,FP))
    FN_ = np.logical_and(FN,FN_)
    FN_ = np.logical_or(FN_, np.logical_xor(gt_ske,np.logical_and(TP,gt_ske))) # preserve the structural components of FN

    ## 2. =============Find exact polygon control points and independent regions==============
    ## find contours from FP_
    ctrs_FP, hier_FP = cv.findContours(FP_.astype(np.uint8), cv.RETR_TREE, cv.CHAIN_APPROX_NONE)
    ## find control points and independent regions for human correction
    bdies_FP, indep_cnt_FP = filter_bdy_cond(ctrs_FP, FP_, np.logical_or(TP,FN_))
    ## find contours from FN_
    ctrs_FN, hier_FN = cv.findContours(FN_.astype(np.uint8), cv.RETR_TREE, cv.CHAIN_APPROX_NONE)
    ## find control points and independent regions for human correction
    bdies_FN, indep_cnt_FN = filter_bdy_cond(ctrs_FN, FN_, 1-np.logical_or(np.logical_or(TP,FP_),FN_))

    poly_FP, poly_FP_len, poly_FP_point_cnt = approximate_RDP(bdies_FP,epsilon=epsilon)
    poly_FN, poly_FN_len, poly_FN_point_cnt = approximate_RDP(bdies_FN,epsilon=epsilon)

    return poly_FP_point_cnt, indep_cnt_FP, poly_FN_point_cnt, indep_cnt_FN

def compute_hce(pred_root,gt_root,gt_ske_root):

    gt_name_list = glob(pred_root+'/*.png')
    gt_name_list = sorted([x.split('/')[-1] for x in gt_name_list])

    hces = []
    for gt_name in tqdm(gt_name_list, total=len(gt_name_list)):
        gt_path = os.path.join(gt_root, gt_name)
        pred_path = os.path.join(pred_root, gt_name)

        gt = cv.imread(gt_path, cv.IMREAD_GRAYSCALE)
        pred = cv.imread(pred_path, cv.IMREAD_GRAYSCALE)

        ske_path = os.path.join(gt_ske_root,gt_name)
        if os.path.exists(ske_path):
            ske = cv.imread(ske_path,cv.IMREAD_GRAYSCALE)
            ske = ske>128
        else:
            ske = skeletonize(gt>128)

        FP_points, FP_indep, FN_points, FN_indep = relax_HCE(gt, pred,ske)
        print(gt_path.split('/')[-1],FP_points, FP_indep, FN_points, FN_indep)
        hces.append([FP_points, FP_indep, FN_points, FN_indep, FP_points+FP_indep+FN_points+FN_indep])

    hce_metric ={'names': gt_name_list,
                 'hces': hces}


    file_metric = open(pred_root+'/hce_metric.pkl','wb')
    pkl.dump(hce_metric,file_metric)
    # file_metrics.write(cmn_metrics)
    file_metric.close()

    return np.mean(np.array(hces)[:,-1])

def main():

    gt_root = "../DIS5K/DIS-VD/gt"
    gt_ske_root = ""
    pred_root = "../Results/isnet(ours)/DIS-VD"

    print("The average HCE metric: ", compute_hce(pred_root,gt_root,gt_ske_root))


if __name__ == '__main__':
    main()