qurator-spk/eynollah
1
0
Fork 0
mirror of https://github.com/qurator-spk/eynollah.git synced 2026-03-02 05:11:57 +01:00
Code Issues Projects Releases Wiki Activity

do_prediction*: simplify (esp. indexing/slicing)

Browse source
This commit is contained in:
Robert Sachunsky 2026-03-01 04:37:20 +01:00
parent 3b56fa2a5b
commit 686f1d34aa
1 changed files with 116 additions and 316 deletions
Download patch file Download diff file Expand all files Collapse all files

388
src/eynollah/eynollah.py
View file

@ -724,25 +724,21 @@ class Eynollah:
if thresholding_for_some_classes_in_light_version: if thresholding_for_some_classes_in_light_version:
seg_not_base = label_p_pred[:,:,:,4] seg_not_base = label_p_pred[:,:,:,4]
seg_not_base[seg_not_base>0.03] =1 seg_not_base = (seg_not_base > 0.03).astype(int)
seg_not_base[seg_not_base<1] =0
seg_line = label_p_pred[:,:,:,3] seg_line = label_p_pred[:,:,:,3]
seg_line[seg_line>0.1] =1 seg_line = (seg_line > 0.1).astype(int)
seg_line[seg_line<1] =0
seg_background = label_p_pred[:,:,:,0] seg_background = label_p_pred[:,:,:,0]
seg_background[seg_background>0.25] =1 seg_background = (seg_background > 0.25).astype(int)
seg_background[seg_background<1] =0
seg[seg_not_base==1]=4 seg[seg_not_base==1]=4
seg[seg_background==1]=0 seg[seg_background==1]=0
seg[(seg_line==1) & (seg==0)]=3 seg[(seg_line==1) & (seg==0)]=3
if thresholding_for_artificial_class_in_light_version: if thresholding_for_artificial_class_in_light_version:
seg_art = label_p_pred[:,:,:,2] seg_art = label_p_pred[:,:,:,2]
seg_art = (seg_art >= threshold_art_class_textline).astype(int)
seg_art[seg_art<threshold_art_class_textline] = 0
seg_art[seg_art>0] =1
##seg[seg_art==1]=2 ##seg[seg_art==1]=2
@ -759,113 +755,51 @@ class Eynollah:
index_x_u_in = list_x_u[indexer_inside_batch] index_x_u_in = list_x_u[indexer_inside_batch]
index_x_d_in = list_x_d[indexer_inside_batch] index_x_d_in = list_x_d[indexer_inside_batch]
if i_batch == 0 and j_batch == 0: where = np.index_exp[index_y_d_in:index_y_u_in,
prediction_true[index_y_d_in + 0:index_y_u_in - margin, index_x_d_in:index_x_u_in]
index_x_d_in + 0:index_x_u_in - margin] = \ if (i_batch == 0 and
seg_in[0:-margin or None, j_batch == 0):
0:-margin or None, inbox = np.index_exp[0:-margin or None,
np.newaxis]
if thresholding_for_artificial_class_in_light_version:
prediction_true[index_y_d_in + 0:index_y_u_in - margin,
index_x_d_in + 0:index_x_u_in - margin, 1] = \
seg_in_art[0:-margin or None,
0:-margin or None] 0:-margin or None]
elif (i_batch == nxf - 1 and
elif i_batch == nxf - 1 and j_batch == nyf - 1: j_batch == nyf - 1):
prediction_true[index_y_d_in + margin:index_y_u_in - 0, inbox = np.index_exp[margin:,
index_x_d_in + margin:index_x_u_in - 0] = \
seg_in[margin:,
margin:,
np.newaxis]
if thresholding_for_artificial_class_in_light_version:
prediction_true[index_y_d_in + margin:index_y_u_in - 0,
index_x_d_in + margin:index_x_u_in - 0, 1] = \
seg_in_art[margin:,
margin:] margin:]
elif (i_batch == 0 and
elif i_batch == 0 and j_batch == nyf - 1: j_batch == nyf - 1):
prediction_true[index_y_d_in + margin:index_y_u_in - 0, inbox = np.index_exp[margin:,
index_x_d_in + 0:index_x_u_in - margin] = \
seg_in[margin:,
0:-margin or None,
np.newaxis]
if thresholding_for_artificial_class_in_light_version:
prediction_true[index_y_d_in + margin:index_y_u_in - 0,
index_x_d_in + 0:index_x_u_in - margin, 1] = \
seg_in_art[margin:,
0:-margin or None] 0:-margin or None]
elif (i_batch == nxf - 1 and
elif i_batch == nxf - 1 and j_batch == 0: j_batch == 0):
prediction_true[index_y_d_in + 0:index_y_u_in - margin, inbox = np.index_exp[0:-margin or None,
index_x_d_in + margin:index_x_u_in - 0] = \
seg_in[0:-margin or None,
margin:,
np.newaxis]
if thresholding_for_artificial_class_in_light_version:
prediction_true[index_y_d_in + 0:index_y_u_in - margin,
index_x_d_in + margin:index_x_u_in - 0, 1] = \
seg_in_art[0:-margin or None,
margin:] margin:]
elif (i_batch == 0 and
elif i_batch == 0 and j_batch != 0 and j_batch != nyf - 1: j_batch != 0 and
prediction_true[index_y_d_in + margin:index_y_u_in - margin, j_batch != nyf - 1):
index_x_d_in + 0:index_x_u_in - margin] = \ inbox = np.index_exp[margin:-margin or None,
seg_in[margin:-margin or None,
0:-margin or None,
np.newaxis]
if thresholding_for_artificial_class_in_light_version:
prediction_true[index_y_d_in + margin:index_y_u_in - margin,
index_x_d_in + 0:index_x_u_in - margin, 1] = \
seg_in_art[margin:-margin or None,
0:-margin or None] 0:-margin or None]
elif (i_batch == nxf - 1 and
elif i_batch == nxf - 1 and j_batch != 0 and j_batch != nyf - 1: j_batch != 0 and
prediction_true[index_y_d_in + margin:index_y_u_in - margin, j_batch != nyf - 1):
index_x_d_in + margin:index_x_u_in - 0] = \ inbox = np.index_exp[margin:-margin or None,
seg_in[margin:-margin or None,
margin:,
np.newaxis]
if thresholding_for_artificial_class_in_light_version:
prediction_true[index_y_d_in + margin:index_y_u_in - margin,
index_x_d_in + margin:index_x_u_in - 0, 1] = \
seg_in_art[margin:-margin or None,
margin:] margin:]
elif (i_batch != 0 and
elif i_batch != 0 and i_batch != nxf - 1 and j_batch == 0: i_batch != nxf - 1 and
prediction_true[index_y_d_in + 0:index_y_u_in - margin, j_batch == 0):
index_x_d_in + margin:index_x_u_in - margin] = \ inbox = np.index_exp[0:-margin or None,
seg_in[0:-margin or None,
margin:-margin or None,
np.newaxis]
if thresholding_for_artificial_class_in_light_version:
prediction_true[index_y_d_in + 0:index_y_u_in - margin,
index_x_d_in + margin:index_x_u_in - margin, 1] = \
seg_in_art[0:-margin or None,
margin:-margin or None] margin:-margin or None]
elif (i_batch != 0 and
elif i_batch != 0 and i_batch != nxf - 1 and j_batch == nyf - 1: i_batch != nxf - 1 and
prediction_true[index_y_d_in + margin:index_y_u_in - 0, j_batch == nyf - 1):
index_x_d_in + margin:index_x_u_in - margin] = \ inbox = np.index_exp[margin:,
seg_in[margin:,
margin:-margin or None,
np.newaxis]
if thresholding_for_artificial_class_in_light_version:
prediction_true[index_y_d_in + margin:index_y_u_in - 0,
index_x_d_in + margin:index_x_u_in - margin, 1] = \
seg_in_art[margin:,
margin:-margin or None] margin:-margin or None]
else: else:
prediction_true[index_y_d_in + margin:index_y_u_in - margin, inbox = np.index_exp[margin:-margin or None,
index_x_d_in + margin:index_x_u_in - margin] = \
seg_in[margin:-margin or None,
margin:-margin or None,
np.newaxis]
if thresholding_for_artificial_class_in_light_version:
prediction_true[index_y_d_in + margin:index_y_u_in - margin,
index_x_d_in + margin:index_x_u_in - margin, 1] = \
seg_in_art[margin:-margin or None,
margin:-margin or None] margin:-margin or None]
prediction_true[where][inbox] = seg_in[inbox + (np.newaxis,)]
if thresholding_for_artificial_class_in_light_version:
prediction_true[where][inbox + (1,)] = seg_in_art[inbox]
indexer_inside_batch += 1 indexer_inside_batch += 1
@ -885,11 +819,7 @@ class Eynollah:
kernel_min = np.ones((3, 3), np.uint8) kernel_min = np.ones((3, 3), np.uint8)
prediction_true[:,:,0][prediction_true[:,:,0]==2] = 0 prediction_true[:,:,0][prediction_true[:,:,0]==2] = 0
skeleton_art = skeletonize(prediction_true[:,:,1]) skeleton_art = skeletonize(prediction_true[:,:,1]).astype(np.uint8)
skeleton_art = skeleton_art*1
skeleton_art = skeleton_art.astype('uint8')
skeleton_art = cv2.dilate(skeleton_art, kernel_min, iterations=1) skeleton_art = cv2.dilate(skeleton_art, kernel_min, iterations=1)
prediction_true[:,:,0][skeleton_art==1]=2 prediction_true[:,:,0][skeleton_art==1]=2
@ -924,18 +854,13 @@ class Eynollah:
if thresholding_for_artificial_class_in_light_version: if thresholding_for_artificial_class_in_light_version:
kernel_min = np.ones((3, 3), np.uint8) kernel_min = np.ones((3, 3), np.uint8)
seg_art = label_p_pred[0,:,:,4] seg_art = label_p_pred[0,:,:,4]
seg_art[seg_art<threshold_art_class_layout] =0 seg_art = (seg_art >= threshold_art_class_layout).astype(int)
seg_art[seg_art>0] =1
#seg[seg_art==1]=4 #seg[seg_art==1]=4
seg_art = resize_image(seg_art, img_h_page, img_w_page).astype(np.uint8) seg_art = resize_image(seg_art, img_h_page, img_w_page).astype(np.uint8)
prediction_true[:,:,0][prediction_true[:,:,0]==4] = 0 prediction_true[:,:,0][prediction_true[:,:,0]==4] = 0
skeleton_art = skeletonize(seg_art) skeleton_art = skeletonize(seg_art).astype(np.uint8)
skeleton_art = skeleton_art*1
skeleton_art = skeleton_art.astype('uint8')
skeleton_art = cv2.dilate(skeleton_art, kernel_min, iterations=1) skeleton_art = cv2.dilate(skeleton_art, kernel_min, iterations=1)
prediction_true[:,:,0][skeleton_art==1] = 4 prediction_true[:,:,0][skeleton_art==1] = 4
@ -948,6 +873,8 @@ class Eynollah:
img = resize_image(img, img.shape[0], img_width_model) img = resize_image(img, img.shape[0], img_width_model)
self.logger.debug("Patch size: %sx%s", img_height_model, img_width_model) self.logger.debug("Patch size: %sx%s", img_height_model, img_width_model)
thresholding = (thresholding_for_artificial_class_in_light_version or
thresholding_for_some_classes_in_light_version)
margin = int(marginal_of_patch_percent * img_height_model) margin = int(marginal_of_patch_percent * img_height_model)
width_mid = img_width_model - 2 * margin width_mid = img_width_model - 2 * margin
height_mid = img_height_model - 2 * margin height_mid = img_height_model - 2 * margin
@ -974,16 +901,8 @@ class Eynollah:
img_patch = np.zeros((n_batch_inference, img_height_model, img_width_model, 3)) img_patch = np.zeros((n_batch_inference, img_height_model, img_width_model, 3))
for i in range(nxf): for i in range(nxf):
for j in range(nyf): for j in range(nyf):
if i == 0:
index_x_d = i * width_mid index_x_d = i * width_mid
index_x_u = index_x_d + img_width_model index_x_u = index_x_d + img_width_model
else:
index_x_d = i * width_mid
index_x_u = index_x_d + img_width_model
if j == 0:
index_y_d = j * height_mid
index_y_u = index_y_d + img_height_model
else:
index_y_d = j * height_mid index_y_d = j * height_mid
index_y_u = index_y_d + img_height_model index_y_u = index_y_d + img_height_model
if index_x_u > img_w: if index_x_u > img_w:
@ -1000,7 +919,8 @@ class Eynollah:
list_y_d.append(index_y_d) list_y_d.append(index_y_d)
list_y_u.append(index_y_u) list_y_u.append(index_y_u)
img_patch[batch_indexer] = img[index_y_d:index_y_u, index_x_d:index_x_u] img_patch[batch_indexer] = img[index_y_d:index_y_u,
index_x_d:index_x_u]
batch_indexer += 1 batch_indexer += 1
if (batch_indexer == n_batch_inference or if (batch_indexer == n_batch_inference or
@ -1012,29 +932,25 @@ class Eynollah:
if thresholding_for_some_classes_in_light_version: if thresholding_for_some_classes_in_light_version:
seg_art = label_p_pred[:,:,:,4] seg_art = label_p_pred[:,:,:,4]
seg_art[seg_art<threshold_art_class_layout] =0
seg_art[seg_art>0] =1 seg_art = (seg_art >= threshold_art_class_layout).astype(int)
seg_line = label_p_pred[:,:,:,3] seg_line = label_p_pred[:,:,:,3]
seg_line[seg_line>0.4] =1#seg_line[seg_line>0.5] =1#seg_line[seg_line>0.1] =1 seg_line = (seg_line > 0.4).astype(int)
seg_line[seg_line<1] =0
##seg[seg_art==1]=4 ##seg[seg_art==1]=4
#seg[(seg_line==1) & (seg==0)]=3 #seg[(seg_line==1) & (seg==0)]=3
if thresholding_for_artificial_class_in_light_version: if thresholding_for_artificial_class_in_light_version:
seg_art = label_p_pred[:,:,:,2] seg_art = label_p_pred[:,:,:,2]
seg_art[seg_art<threshold_art_class_textline] = 0 seg_art = (seg_art >= threshold_art_class_textline).astype(int)
seg_art[seg_art>0] =1
##seg[seg_art==1]=2 ##seg[seg_art==1]=2
indexer_inside_batch = 0 indexer_inside_batch = 0
for i_batch, j_batch in zip(list_i_s, list_j_s): for i_batch, j_batch in zip(list_i_s, list_j_s):
seg_in = seg[indexer_inside_batch] seg_in = seg[indexer_inside_batch]
if (thresholding_for_artificial_class_in_light_version or if thresholding:
thresholding_for_some_classes_in_light_version):
seg_in_art = seg_art[indexer_inside_batch] seg_in_art = seg_art[indexer_inside_batch]
index_y_u_in = list_y_u[indexer_inside_batch] index_y_u_in = list_y_u[indexer_inside_batch]
@ -1043,164 +959,56 @@ class Eynollah:
index_x_u_in = list_x_u[indexer_inside_batch] index_x_u_in = list_x_u[indexer_inside_batch]
index_x_d_in = list_x_d[indexer_inside_batch] index_x_d_in = list_x_d[indexer_inside_batch]
if i_batch == 0 and j_batch == 0: where = np.index_exp[index_y_d_in:index_y_u_in,
prediction_true[index_y_d_in + 0:index_y_u_in - margin, index_x_d_in:index_x_u_in]
index_x_d_in + 0:index_x_u_in - margin] = \ if (i_batch == 0 and
seg_in[0:-margin or None, j_batch == 0):
0:-margin or None, inbox = np.index_exp[0:-margin or None,
np.newaxis]
confidence_matrix[index_y_d_in + 0:index_y_u_in - margin,
index_x_d_in + 0:index_x_u_in - margin] = \
label_p_pred[0, 0:-margin or None,
0:-margin or None,
1]
if (thresholding_for_artificial_class_in_light_version or
thresholding_for_some_classes_in_light_version):
prediction_true[index_y_d_in + 0:index_y_u_in - margin,
index_x_d_in + 0:index_x_u_in - margin, 1] = \
seg_in_art[0:-margin or None,
0:-margin or None] 0:-margin or None]
elif (i_batch == nxf - 1 and
elif i_batch == nxf - 1 and j_batch == nyf - 1: j_batch == nyf - 1):
prediction_true[index_y_d_in + margin:index_y_u_in - 0, inbox = np.index_exp[margin:,
index_x_d_in + margin:index_x_u_in - 0] = \
seg_in[margin:,
margin:,
np.newaxis]
confidence_matrix[index_y_d_in + margin:index_y_u_in - 0,
index_x_d_in + margin:index_x_u_in - 0] = \
label_p_pred[0, margin:,
margin:,
1]
if (thresholding_for_artificial_class_in_light_version or
thresholding_for_some_classes_in_light_version):
prediction_true[index_y_d_in + margin:index_y_u_in - 0,
index_x_d_in + margin:index_x_u_in - 0, 1] = \
seg_in_art[margin:,
margin:] margin:]
elif (i_batch == 0 and
elif i_batch == 0 and j_batch == nyf - 1: j_batch == nyf - 1):
prediction_true[index_y_d_in + margin:index_y_u_in - 0, inbox = np.index_exp[margin:,
index_x_d_in + 0:index_x_u_in - margin] = \
seg_in[margin:,
0:-margin or None,
np.newaxis]
confidence_matrix[index_y_d_in + margin:index_y_u_in - 0,
index_x_d_in + 0:index_x_u_in - margin] = \
label_p_pred[0, margin:,
0:-margin or None,
1]
if (thresholding_for_artificial_class_in_light_version or
thresholding_for_some_classes_in_light_version):
prediction_true[index_y_d_in + margin:index_y_u_in - 0,
index_x_d_in + 0:index_x_u_in - margin, 1] = \
seg_in_art[margin:,
0:-margin or None] 0:-margin or None]
elif (i_batch == nxf - 1 and
elif i_batch == nxf - 1 and j_batch == 0: j_batch == 0):
prediction_true[index_y_d_in + 0:index_y_u_in - margin, inbox = np.index_exp[0:-margin or None,
index_x_d_in + margin:index_x_u_in - 0] = \
seg_in[0:-margin or None,
margin:,
np.newaxis]
confidence_matrix[index_y_d_in + 0:index_y_u_in - margin,
index_x_d_in + margin:index_x_u_in - 0] = \
label_p_pred[0, 0:-margin or None,
margin:,
1]
if (thresholding_for_artificial_class_in_light_version or
thresholding_for_some_classes_in_light_version):
prediction_true[index_y_d_in + 0:index_y_u_in - margin,
index_x_d_in + margin:index_x_u_in - 0, 1] = \
seg_in_art[0:-margin or None,
margin:] margin:]
elif (i_batch == 0 and
elif i_batch == 0 and j_batch != 0 and j_batch != nyf - 1: j_batch != 0 and
prediction_true[index_y_d_in + margin:index_y_u_in - margin, j_batch != nyf - 1):
index_x_d_in + 0:index_x_u_in - margin] = \ inbox = np.index_exp[margin:-margin or None,
seg_in[margin:-margin or None,
0:-margin or None,
np.newaxis]
confidence_matrix[index_y_d_in + margin:index_y_u_in - margin,
index_x_d_in + 0:index_x_u_in - margin] = \
label_p_pred[0, margin:-margin or None,
0:-margin or None,
1]
if (thresholding_for_artificial_class_in_light_version or
thresholding_for_some_classes_in_light_version):
prediction_true[index_y_d_in + margin:index_y_u_in - margin,
index_x_d_in + 0:index_x_u_in - margin, 1] = \
seg_in_art[margin:-margin or None,
0:-margin or None] 0:-margin or None]
elif i_batch == nxf - 1 and j_batch != 0 and j_batch != nyf - 1: elif (i_batch == nxf - 1 and
prediction_true[index_y_d_in + margin:index_y_u_in - margin, j_batch != 0 and
index_x_d_in + margin:index_x_u_in - 0] = \ j_batch != nyf - 1):
seg_in[margin:-margin or None, inbox = np.index_exp[margin:-margin or None,
margin:,
np.newaxis]
confidence_matrix[index_y_d_in + margin:index_y_u_in - margin,
index_x_d_in + margin:index_x_u_in - 0] = \
label_p_pred[0, margin:-margin or None,
margin:,
1]
if (thresholding_for_artificial_class_in_light_version or
thresholding_for_some_classes_in_light_version):
prediction_true[index_y_d_in + margin:index_y_u_in - margin,
index_x_d_in + margin:index_x_u_in - 0, 1] = \
seg_in_art[margin:-margin or None,
margin:] margin:]
elif i_batch != 0 and i_batch != nxf - 1 and j_batch == 0: elif (i_batch != 0 and
prediction_true[index_y_d_in + 0:index_y_u_in - margin, i_batch != nxf - 1 and
index_x_d_in + margin:index_x_u_in - margin] = \ j_batch == 0):
seg_in[0:-margin or None, inbox = np.index_exp[0:-margin or None,
margin:-margin or None,
np.newaxis]
confidence_matrix[index_y_d_in + 0:index_y_u_in - margin,
index_x_d_in + margin:index_x_u_in - margin] = \
label_p_pred[0, 0:-margin or None,
margin:-margin or None,
1]
if (thresholding_for_artificial_class_in_light_version or
thresholding_for_some_classes_in_light_version):
prediction_true[index_y_d_in + 0:index_y_u_in - margin,
index_x_d_in + margin:index_x_u_in - margin, 1] = \
seg_in_art[0:-margin or None,
margin:-margin or None] margin:-margin or None]
elif i_batch != 0 and i_batch != nxf - 1 and j_batch == nyf - 1: elif (i_batch != 0 and
prediction_true[index_y_d_in + margin:index_y_u_in - 0, i_batch != nxf - 1 and
index_x_d_in + margin:index_x_u_in - margin] = \ j_batch == nyf - 1):
seg_in[margin:, inbox = np.index_exp[margin:,
margin:-margin or None,
np.newaxis]
confidence_matrix[index_y_d_in + margin:index_y_u_in - 0,
index_x_d_in + margin:index_x_u_in - margin] = \
label_p_pred[0, margin:,
margin:-margin or None,
1]
if (thresholding_for_artificial_class_in_light_version or
thresholding_for_some_classes_in_light_version):
prediction_true[index_y_d_in + margin:index_y_u_in - 0,
index_x_d_in + margin:index_x_u_in - margin, 1] = \
seg_in_art[margin:,
margin:-margin or None] margin:-margin or None]
else: else:
prediction_true[index_y_d_in + margin:index_y_u_in - margin, inbox = np.index_exp[margin:-margin or None,
index_x_d_in + margin:index_x_u_in - margin] = \
seg_in[margin:-margin or None,
margin:-margin or None,
np.newaxis]
confidence_matrix[index_y_d_in + margin:index_y_u_in - margin,
index_x_d_in + margin:index_x_u_in - margin] = \
label_p_pred[0, margin:-margin or None,
margin:-margin or None,
1]
if (thresholding_for_artificial_class_in_light_version or
thresholding_for_some_classes_in_light_version):
prediction_true[index_y_d_in + margin:index_y_u_in - margin,
index_x_d_in + margin:index_x_u_in - margin, 1] = \
seg_in_art[margin:-margin or None,
margin:-margin or None] margin:-margin or None]
prediction_true[where][inbox] = seg_in[inbox + (np.newaxis,)]
confidence_matrix[where][inbox] = label_p_pred[(0,) + inbox + (1,)]
# rs: why is prediction_true 3ch when only 1st gets used?
# artificial boundary class map should be extra array
# rs: why does confidence_matrix only get text-label scores?
# should be scores at final argmax
if thresholding:
prediction_true[where][inbox + (1,)] = seg_in_art[inbox]
indexer_inside_batch += 1 indexer_inside_batch += 1
list_i_s = [] list_i_s = []
@ -1219,11 +1027,7 @@ class Eynollah:
kernel_min = np.ones((3, 3), np.uint8) kernel_min = np.ones((3, 3), np.uint8)
prediction_true[:,:,0][prediction_true[:,:,0]==2] = 0 prediction_true[:,:,0][prediction_true[:,:,0]==2] = 0
skeleton_art = skeletonize(prediction_true[:,:,1]) skeleton_art = skeletonize(prediction_true[:,:,1]).astype(np.uint8)
skeleton_art = skeleton_art*1
skeleton_art = skeleton_art.astype('uint8')
skeleton_art = cv2.dilate(skeleton_art, kernel_min, iterations=1) skeleton_art = cv2.dilate(skeleton_art, kernel_min, iterations=1)
prediction_true[:,:,0][skeleton_art==1]=2 prediction_true[:,:,0][skeleton_art==1]=2
@ -1232,11 +1036,7 @@ class Eynollah:
kernel_min = np.ones((3, 3), np.uint8) kernel_min = np.ones((3, 3), np.uint8)
prediction_true[:,:,0][prediction_true[:,:,0]==4] = 0 prediction_true[:,:,0][prediction_true[:,:,0]==4] = 0
skeleton_art = skeletonize(prediction_true[:,:,1]) skeleton_art = skeletonize(prediction_true[:,:,1]).astype(np.uint8)
skeleton_art = skeleton_art*1
skeleton_art = skeleton_art.astype('uint8')
skeleton_art = cv2.dilate(skeleton_art, kernel_min, iterations=1) skeleton_art = cv2.dilate(skeleton_art, kernel_min, iterations=1)
prediction_true[:,:,0][skeleton_art==1]=4 prediction_true[:,:,0][skeleton_art==1]=4