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
98
99
100
101
102
103
|
import onnxruntime
import numpy as np
import torch
import torch.nn.functional as F
from PIL import Image
import cv2
from torchvision import transforms
"""
使用GPU进行推理
"""
# GPU推理引擎（使用CPU：providers = ['CPUExecutionProvider']）
providers = ['CUDAExecutionProvider', 'CPUExecutionProvider']
ort_session = onnxruntime.InferenceSession('FastSegFormer_P_224.onnx', providers=providers)
# 读取图像（打开文件的形式）
ori_image = cv2.imdecode(np.fromfile('images/img236.jpg', np.uint8), cv2.IMREAD_COLOR)
# 图像预处理函数
test_transform = transforms.Compose([transforms.Resize(224),
transforms.Normalize(
mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
])
input_image = test_transform(ori_image)
def blend_images(old_image, new_image, alpha):
"""
使用cv2.addWeighted()函数混合两个图像
"""
blended_image = cv2.addWeighted(old_image, alpha, new_image, 1 - alpha, 0)
return blended_image
# 图像分割检测图片方法
def detect_image(model, image, name_classes = None, num_classes = 21, count = False, input_shape = [224, 224], device = 'cpu', weight_type = None):
# 转化为彩色图像
image = cvtColor(image)
image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
# 对输入图像做一个备份
old_img = copy.deepcopy(image)
original_h = np.array(image).shape[0]
original_w = np.array(image).shape[1]
image_data = cv2.resize(image, input_shape, interpolation=cv2.INTER_LINEAR)
# 添加Batch维度
image_data = np.expand_dims(np.transpose(preprocess_input(np.array(image_data, np.float32)), (2, 0, 1)), 0)
if weight_type == 'pth':
with torch.no_grad():
# 转化为张量
images = torch.from_numpy(image_data)
images = images.to(device)
pred = model(images)[0]
pred = F.softmax(pred.permute(1,2,0),dim = -1).cpu().numpy()
pred = cv2.resize(pred, (original_w, original_h), interpolation = cv2.INTER_LINEAR)
pred = pred.argmax(axis=-1)
elif weight_type == 'onnx':
ort_inputs = {'input': image_data}
pred = model.run(['output'], ort_inputs)[0]
pred = pred[0]
# 转化为张量
pred = torch.tensor(pred)
pred = F.softmax(pred.permute(1,2,0),dim = -1).cpu().numpy()
pred = cv2.resize(pred, (original_w, original_h), interpolation = cv2.INTER_LINEAR)
pred = pred.argmax(axis=-1)
if count:
classes_nums = np.zeros([num_classes])
total_points_num = original_h * original_w
print('-' * 63)
print("|%25s | %15s | %15s|"%("Key", "Value", "Ratio"))
print('-' * 63)
for i in range(num_classes):
num = np.sum(pred == i)
ratio = num / total_points_num * 100
if num > 0:
print("|%25s | %15s | %14.2f%%|"%(str(name_classes[i]), str(num), ratio))
print('-' * 63)
classes_nums[i] = num
print("classes_nums:", classes_nums)
if num_classes <= 21:
# 要画的像素颜色（这里使用的是VOC格式的像素颜色）
colors = [ (0, 0, 0), (128, 0, 0), (0, 128, 0), (128, 128, 0), (0, 0, 128), (128, 0, 128), (0, 128, 128),
(128, 128, 128), (64, 0, 0), (192, 0, 0), (64, 128, 0), (192, 128, 0), (64, 0, 128), (192, 0, 128),
(64, 128, 128), (192, 128, 128), (0, 64, 0), (128, 64, 0), (0, 192, 0), (128, 192, 0), (0, 64, 128),
(128, 64, 12)]
seg_img = np.reshape(np.array(colors, np.uint8)[np.reshape(pred, [-1])], [original_h, original_w, -1])
image = blend_images(old_image=old_img, new_image=seg_img, alpha=0.6)
seg_img = cv2.cvtColor(seg_img, cv2.COLOR_RGB2BGR)
image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
# 分别返回预测结果和预测与原图混合的结果
return seg_img, image
# 进行图像分割
result, image_det = detect_image(model=ort_session, image = ori_image,\\
name_classes=["background", "sunburn", "Ulcer", "wind scarring"], num_classes=4,\\
input_shape=[224, 224], device='cuda', weight_type='onnx')
|
