mirror of
https://github.com/xuebinqin/DIS.git
synced 2024-11-29 18:23:08 +01:00
189 lines
6.8 KiB
Python
189 lines
6.8 KiB
Python
## 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()
|