Unet.md 3.58 KB
Newer Older
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
# 图像分割

本示例主要通过Unet模型说明如何使用MIGraphX Python API进行图像分割模型的推理,包括预处理、模型推理。

## 模型简介

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

<img src="../Images/Unet_01.png" style="zoom:80%;" align=middle>

## 预处理

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

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

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

19
​      3.归一化,将数据归一化到[0.0, 1.0]之间
20
21
22
23
24
25

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

```python
def Preprocessing(pil_img, newW, newH):
    assert newW > 0 and newH > 0, 'Scale is too small'  
26
27
28
29
30
31
32
33
34
35
36
37
    img_nd = cv2.cvtColor(pil_img, cv2.COLOR_BGR2RGB)    # BGR转换为RGB
    img_nd = cv2.resize(img_nd, (newW, newH))            # 将图像尺寸修改为256x256
    
    if len(img_nd.shape) == 2:
        img_nd = np.expand_dims(img_nd, axis=2)

    img_trans = img_nd.transpose((2, 0, 1))              # HWC转换为CHW
    img_trans = np.expand_dims(img_trans, 0)             # CHW扩展为NCHW
    img_trans = np.ascontiguousarray(img_trans)          # 保证内存连续存储
    img_trans = img_trans.astype(np.float32)             # 转换成浮点型数据
    if img_trans.max() > 1:                             
        img = img_trans / 255.0                          # 保证数据处于0-1之间的浮点数
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
    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文件是通过使用具有普通背景的汽车图像来训练的。因此,“现实世界“图像的分割结果不完美是意料之中的。为了获得更好的结果,建议对现实世界示例数据集上的模型进行微调。