File size: 3,255 Bytes
edcf5ee |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 |
"""
对测试集进行处理，将图片中的实例遮住，实例部分用原图像素均值填充，生成新的测试集
ver： Feb 21th
"""
import numpy as np
import cv2
import os
def find_all_files(root, suffix=None):
"""
Return a list of file paths ended with specific suffix
"""
res = []
if type(suffix) is str or suffix is None:
for root, _, files in os.walk(root):
for f in files:
if suffix is not None and not f.endswith(suffix):
continue
res.append(os.path.join(root, f))
return res
else:
print('type of suffix is not legal :', type(suffix))
return -1
if __name__ == '__main__':
# 只需要修改数据路径和result路径，new_test与test平级
# 导入测试集image和mask
data_path = 'E:/Study/code/datasets/SIPaKMeD_MIL/test/data/' # MIL数据集的路径
result_path = 'E:/Study/code/datasets/SIPaKMeD_MIL/new_test/'
if not os.path.exists(result_path):
os.makedirs(result_path)
suffix = '.jpg'
# 获取类别名，制作label和类别名的对应字典
class_names = [filename for filename in os.listdir(data_path)
if os.path.isdir(os.path.join(data_path, filename))]
class_names.sort()
cls_idxs = [i for i in range(len(class_names))]
class_id_dict = dict(zip(class_names, cls_idxs))
input_ids = sorted(find_all_files(data_path, suffix=suffix))
# 制作结果路径
for class_name in class_names:
res_path = result_path + class_name
if not os.path.exists(res_path):
os.makedirs(res_path)
for i in range(len(input_ids)):
image_path = input_ids[i]
# 读取image和mask
# CV2 0-255 hwc，in totensor step it will be transformed to chw. ps:PIL(0-1 hwc)
image = np.array(cv2.imread(image_path))
# mask_path is replace the last 'data' by 'mask'
mask_path = "data".join(image_path.split("data")[:-1]) + 'mask' + "".join(image_path.split("data")[-1:])
# mask: 0/255 cv2 hwc
mask = np.array(cv2.imread(mask_path))
mask_norm = np.where(mask > 50, 0, 1)
# new_image_path is replace the last 'test/data' by 'new_test'
new_image_path = "data".join(image_path.split("test/data")[:-1]) + \
'new_test' + "".join(image_path.split("test/data")[-1:])
new_image = image * mask_norm
# 把抠掉的部分填充成原图的像素均值
value_mean_r = int(np.mean(image[:, :, 0]))
value_mean_g = int(np.mean(image[:, :, 1]))
value_mean_b = int(np.mean(image[:, :, 2]))
new_image[:, :, 0][new_image[:, :, 0] == 0] = value_mean_r
new_image[:, :, 1][new_image[:, :, 1] == 0] = value_mean_g
new_image[:, :, 2][new_image[:, :, 2] == 0] = value_mean_b
new_image = new_image.astype(np.uint8)
# # 显示原图，mask，new_image
# images = np.hstack([image, mask, new_image])
# cv2.imshow('Before and after mask', images)
# cv2.waitKey(0)
# 存储新的图片
cv2.imwrite(new_image_path, new_image)
|