Unet.md 3.77 KB
Newer Older
lijian6's avatar
lijian6 committed
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
# 图像分割

本示例采用了经典的Unet模型进行图像分割,模型下载地址:https://www.dropbox.com/s/3ntkhyk30x05uuv/unet_13_256.onnx,将unet_13_256.onnx文件保存在Resource\Models\Segmentation文件夹下。模型结构如下图所示(可以通过netron工具(https://netron.app/)查看模型结构),该模型的输入shape为[batch_size,3,256,256],输出shape为[batch_size,1,256,256],数据排布为NCHW。

![Unet_Image_1](../Images/Unet_Image_1.png)

## 预处理

在将数据输入到模型之前,需要对图像做如下预处理操作:

​      1.尺度变换,将图像resize到256x256大小

​      2.归一化,将数据归一化到[0.0, 1.0]之间

​      3.数据排布,将数据从HWC转换为CHW,再将维度转换为NCHW

本示例代码通过如下方式实现预处理操作:

```python
def Preprocessing(pil_img, newW, newH):
    assert newW > 0 and newH > 0, 'Scale is too small'  # 判断规定的尺寸是否为空
    img_nd = cv2.resize(pil_img, (newW, newH))          # 将图像尺寸修改为256x256
    img_nd = cv2.cvtColor(img_nd, cv2.COLOR_BGR2RGB)    # BGR转换为RGB

    img_trans = img_nd.transpose((2, 0, 1))             # HWC转换为CHW
    if img_trans.max() > 1:                             # 保证数据处于0-1之间的浮点数
        img_trans = img_trans / 255.0  
    img_trans = np.expand_dims(img_trans, 0)            # CHW转换为NCHW
    img = img_trans.astype(np.float32)                  # 转换成浮点型数据
    return img
```

## 推理

完成图像预处理后,就可以执行推理,得到推理结果。

```python
    # 加载模型
    model = migraphx.parse_onnx("../../Resource/Models/Segmentation/unet_13_256.onnx")

    # 编译模型
    model.compile(migraphx.get_target("gpu"), device_id=0)   #device_id: 设置GPU设备,默认为0号设备

    # 图像预处理
    img = cv2.imread("../../Resource/Images/car1.jpeg")
    input_img = Preprocessing(img, 256, 256)

    # 模型推理
    mask = model.run({'inputs':input_img})      
    output_mask = np.array(mask[0])[0]          # 获取推理结果,shape为(1,256,256)
    probs = Sigmoid(output_mask)                # 计算sigmoid值

    # 0/1像素值,当概率值大于0.996时,值为255,小于等于0.996时,值为0
    output_mask[probs > 0.996] = 255            
    output_mask[probs <= 0.996] = 0

    output = output_mask.astype(np.uint8)[0]  # 将浮点数类型转换为无符号整型,shape为(256,256)
    cv2.imwrite("output.jpg", output)         # 保存图像分割结果

```

1.Preprocessing函数返回预处理后的数据(numpy类型),然后通过model.run({'inputs':input_img})得到推理结果,因为只有输入一张图片,所以通过mask[0]获取第一个输出节点的数据即可。

2.模型得到的推理结果并不能直接作为分割结果。首先,需要计算sigmoid值,当概率值大于0.996时值为255,小于等于0.996时值为0。其次,将数据类型转换为无符号整形。最终,保存结果得到分割图像。

注:本次采用的模型权重onnx文件是通过使用具有普通背景的汽车图像来训练的。因此,“现实世界“图像的分割结果不完美是意料之中的。为了获得更好的结果,建议对现实世界示例数据集上的模型进行微调。

## 运行示例

1.参考《MIGraphX教程》中的安装方法安装MIGraphX并设置好PYTHONPATH

2.安装依赖:

```Python
# 进入migraphx samples工程根目录
cd <path_to_migraphx_samples> 

# 进入示例程序目录
cd Python/Segmentation

# 安装依赖
pip install -r requirements.txt
```

3.在Python/Segmentation目录下执行如下命令运行该示例程序:

```
python Unet.py
```

输出结果为:

![Unet_output_python](../Images/Unet_output_python.jpg)