115 lines
3.1 KiB
Python
115 lines
3.1 KiB
Python
import os
|
||
|
||
import numpy as np
|
||
import yaml
|
||
from PIL import Image
|
||
import data_utils
|
||
import torch
|
||
from pathlib import Path
|
||
import model_utils
|
||
import utils
|
||
from matplotlib import pyplot as plt
|
||
|
||
|
||
|
||
"""
|
||
预测
|
||
@:param net: 网络模型
|
||
@:param image: 图像
|
||
@:param cls_name: 类别名
|
||
@:param predict_config: 预测配置文件路径
|
||
@:param model_config: 模型配置文件路径
|
||
|
||
@:return mask: [image_height, image_width]，元素类型为bool
|
||
"""
|
||
def predict(
|
||
net,
|
||
image: Image,
|
||
cls_name,
|
||
predict_config=Path("config") / "predict.yaml",
|
||
model_config=Path("config") / "model.yaml"
|
||
):
|
||
with model_config.open("r", encoding="utf-8") as mcf:
|
||
model_config = yaml.load(mcf, Loader=yaml.FullLoader)
|
||
classes = model_config["classes"]
|
||
|
||
with predict_config.open("r", encoding="utf-8") as pcf:
|
||
predict_config = yaml.load(pcf, yaml.FullLoader)
|
||
device = predict_config["device"]
|
||
image = data_utils.pil2tensor(image, device)
|
||
if len(image.shape) == 3:
|
||
image = torch.unsqueeze(image, dim=0)
|
||
batch_size, _, image_height, image_width = image.shape
|
||
|
||
prediction = data_utils.inv_one_hot_of_outputs(
|
||
torch.transpose(
|
||
net(image),
|
||
-2,
|
||
-1
|
||
).reshape(batch_size, len(classes), image_height, image_width),
|
||
device
|
||
)
|
||
|
||
mask = torch.squeeze(
|
||
prediction == utils.get_label_of_cls(classes, cls_name)[0]
|
||
)
|
||
|
||
return mask
|
||
|
||
|
||
"""
|
||
将预测结果与原图混合
|
||
|
||
@:param net: 神经网络模型
|
||
@:param image: 原图
|
||
@:param mask: predict的对应某一类别的mask
|
||
@:param mask: 神经网络的预测结果
|
||
@:param classes: 所有类别
|
||
@:param cls_name: 类别
|
||
@:param colors: 所有类别对应的颜色列表
|
||
@:return 混合后的图像
|
||
"""
|
||
def blend(
|
||
image: Image,
|
||
mask,
|
||
classes,
|
||
cls_name,
|
||
colors
|
||
):
|
||
mask = mask.to(device="cpu").numpy()
|
||
new_image = np.zeros((*mask.shape, 3), dtype=np.uint8)
|
||
new_image[mask] = utils.get_color_of_cls(classes, colors, cls_name)
|
||
new_image = Image.fromarray(new_image)
|
||
blend_image = Image.blend(image, new_image, 0.5)
|
||
return blend_image
|
||
|
||
|
||
|
||
"""
|
||
展示图像
|
||
@:param 需要进行展示的图像，图像尺寸应为[height, width, channels=3]
|
||
"""
|
||
def show_image(image):
|
||
plt.imshow(image)
|
||
plt.show()
|
||
|
||
|
||
|
||
|
||
if __name__ == "__main__":
|
||
with Path(os.path.sep.join(["config", "model.yaml"])).open("r", encoding="utf-8") as f:
|
||
model_config = yaml.load(f, Loader=yaml.FullLoader)
|
||
classes = model_config["classes"]
|
||
|
||
colors = utils.get_colors(len(classes))
|
||
|
||
image_path = os.path.sep.join([
|
||
"dataset", "test", "biomass_image_train_0233_8.jpg"
|
||
])
|
||
|
||
cls_name = "leaf"
|
||
net = model_utils.get_model(False)
|
||
model_utils.init_model(False, net)
|
||
image = Image.open(image_path)
|
||
mask = predict(net, image, cls_name)
|
||
show_image(blend(image, mask, classes, cls_name, colors)) |