|
import cv2 |
|
import numpy as np |
|
|
|
|
|
def erosion_hair_region(img_c1u, num_iter): |
|
|
|
fil = np.array([[ 0, 0.25, 0], |
|
[ 0.25, -1.0, 0.25], |
|
[ 0, 0.25, 0]]) |
|
|
|
ddepth = cv2.CV_32FC1 |
|
temp_img = img_c1u.copy() |
|
|
|
temp_img[temp_img == 1] = 3 |
|
temp_img[temp_img == 2] = 1 |
|
temp_img[temp_img == 3] = 2 |
|
|
|
|
|
|
|
img_f = temp_img.astype(np.float32) |
|
|
|
for _ in range(num_iter): |
|
|
|
img_res = cv2.filter2D(img_f, ddepth, fil, borderType=cv2.BORDER_CONSTANT) |
|
mask_reion = (img_c1u == 2) * (img_res < -0.01) |
|
|
|
img_f[mask_reion] = 0.0 |
|
|
|
|
|
|
|
|
|
|
|
|
|
res = img_f.astype(np.uint8) |
|
res[res == 1] = 3 |
|
res[res == 2] = 1 |
|
res[res == 3] = 2 |
|
|
|
return res |
|
|
|
|
|
def extract_max_region(label_img, tar_value): |
|
mask_img = np.zeros_like(label_img) |
|
mask_img[label_img == tar_value] = 1 |
|
num_labels, label_img = cv2.connectedComponents(mask_img, connectivity=8) |
|
|
|
max_label = -1 |
|
max_area = -1.0 |
|
|
|
for i in range(1, num_labels): |
|
cur_area = np.sum(label_img == i) |
|
if cur_area > max_area: |
|
max_label = i |
|
max_area = cur_area |
|
|
|
label_img[label_img != max_label] = 0 |
|
label_img[label_img == max_label] = 255 |
|
return label_img |
|
|
|
|
|
def remover_free_block(img_c1u): |
|
temp_img = img_c1u.copy() |
|
|
|
temp_img[temp_img > 0.5] = 1 |
|
label_img = extract_max_region(temp_img, 1) |
|
|
|
img_c1u[label_img != 255] = 0 |
|
|
|
return img_c1u |
|
|
|
|
|
def correct_hair_mask(mask_img): |
|
|
|
mask_img = remover_free_block(mask_img) |
|
mask_img = erosion_hair_region(mask_img, 7) |
|
mask_img = remover_free_block(mask_img) |
|
|
|
return mask_img |
