-
Notifications
You must be signed in to change notification settings - Fork 0
/
Predict.py
97 lines (86 loc) · 4.8 KB
/
Predict.py
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
86
87
88
89
90
91
92
93
94
95
96
97
from Loss import dice_loss,dice_coefficient,iou
from Preprocess import preprocess_image_oneband,preprocess_image_twoband
from tensorflow.keras.preprocessing.image import img_to_array, load_img
from tensorflow.keras.models import load_model
import numpy as np
from skimage.transform import resize
import os
import skimage.io as io
import zipfile
def main():
# 加载模型时提供自定义损失函数的定义
model_name='unet_band1_val_best.hdf5'
# 指定要预测的图片目录和保存掩码的目录
predict_dir = 'data/filter/image/band_1'
save_dir = 'data/test_band_1'
#指定输出的mask的阈值
mask_value=0.5
#输入网络的图片大小
target_size=(256, 256)
#原始预测图片大小
original_size=(512, 512)
# 指定目标目录和压缩文件名
target_directory = "data/test_result"
zip_file_name = "data/test_result/dhdata.zip"
#加载模型预测图片以及打包
model = load_model(model_name,
custom_objects={'dice_loss': dice_loss, 'dice_coefficient': dice_coefficient, 'iou': iou})
if not os.path.exists(save_dir):
os.makedirs(save_dir)
# 遍历目录中的所有图片，对每张图片进行预测并保存掩码
for image_file in os.listdir(predict_dir):
image_path = os.path.join(predict_dir, image_file)
#可以增加mode，使快速区分oneband以及多band
predict_oneband_save(model, image_path, target_size=target_size, original_size=original_size, mask_value=mask_value, save_path=save_dir)
#predict_twoband_save(model, image_path, target_size=target_size, original_size=original_size, mask_value=mask_value, save_path=save_dir)
print("预测完成，掩码已保存至：", save_dir)
# 调用函数进行压缩
zip_images(target_directory, zip_file_name)
def load_oneband_image(image_path, target_size=(256, 256)):
"""预处理图片：读取、缩放和归一化"""
image = load_img(image_path, color_mode='grayscale', target_size=target_size)
image = img_to_array(image)
preprocess_image_oneband(image)
image = np.expand_dims(image, axis=0) # 增加批次维度
return image
def load_twoband_image(image_path, target_size=(256, 256)):
image=preprocess_image_twoband(image_path)
"""针对双波段tif影像的预处理函数"""
image = np.expand_dims(image, axis=0) # 增加批次维度
return image
def predict_oneband_save(model, image_path, target_size=(256, 256), original_size=(512, 512),mask_value=0.5,
save_path='data/predicted_masks'):
"""对单张图片进行预测，并保存掩码。现在会将预测的掩码调整为原始尺寸。"""
image = load_oneband_image(image_path, target_size)
mask = model.predict(image)[0] # 预测并获取掩码
mask = (mask > mask_value ).astype(np.float32) # 二值化处理，这里改用 float32 以防在 resize 过程中出现问题
# 将掩码调整为原始尺寸
mask_resized = resize(mask, original_size, mode='constant', preserve_range=True)
mask_resized = (mask_resized > mask_value).astype(np.uint8) * 255 # 再次二值化处理，确保掩码是黑白的
file_name = os.path.basename(image_path)
# 去掉文件名中的 '_band_1' 词缀，并确保保存为PNG格式
file_name_without_band = file_name.replace('_band_1', '')
mask_save_path = os.path.join(save_path, os.path.splitext(file_name_without_band)[0] + '_msk.png')
io.imsave(mask_save_path, mask_resized.squeeze()) # 保存掩码图片
def predict_twoband_save(model, image_path, target_size=(256, 256), original_size=(512, 512),mask_value=0.5,
save_path='data/predicted_masks'):
"""对单张双波段图片进行预测，并保存掩码。现在会将预测的掩码调整为原始尺寸。"""
image = load_twoband_image(image_path, target_size)
mask = model.predict(image)[0] # 预测并获取掩码
mask = (mask > mask_value).astype(np.float32) # 二值化处理
# 将掩码调整为原始尺寸
mask_resized = resize(mask, original_size, mode='constant', preserve_range=True)
mask_resized = (mask_resized > mask_value).astype(np.uint8) * 255 # 再次二值化处理
file_name = os.path.basename(image_path)
mask_save_path = os.path.join(save_path, os.path.splitext(file_name)[0] + '_msk.png')
io.imsave(mask_save_path, mask_resized.squeeze()) # 保存掩码图片
def zip_images(directory, zip_file_name):
# 获取指定目录下所有文件的路径
file_paths = [os.path.join(directory, f) for f in os.listdir(directory)]
# 初始化 zip 文件
with zipfile.ZipFile(zip_file_name, 'w') as zipf:
# 将目录下的所有文件添加到 zip 文件中
for file_path in file_paths:
zipf.write(file_path, os.path.basename(file_path))
if __name__ == "__main__":
main()