IOS 上的图像分割 DEEPLABV3-技术圈

介绍

语义图像分割是一项计算机视觉任务，它使用语义标签来标记输入图像的特定区域。PyTorch 语义图像分割DeepLabV3 模型可用于标记具有20 个语义类的图像区域，例如，自行车、公共汽车、汽车、狗和人。图像分割模型在自动驾驶和场景理解等应用中非常有用。

在本教程中，我们将提供有关如何在 iOS 上准备和运行 PyTorch DeepLabV3 模型的分步指南，带您从开始拥有一个您可能想要在 iOS 上使用的模型到最终拥有一个完整的模型。使用该模型的 iOS 应用程序。我们还将介绍有关如何检查下一个最喜欢的预训练 PyTorch 模型是否可以在 iOS 上运行以及如何避免陷阱的实用和一般技巧。

笔记

在学习本教程之前，您应该查看适用于 iOS 的 PyTorch Mobile，并快速尝试一下PyTorch iOS HelloWorld示例应用程序。本教程将超越图像分类模型，通常是部署在移动设备上的第一种模型。本教程的完整代码存储库可在此处获得。

学习目标

在本教程中，您将学习如何：

为 iOS 部署转换 DeepLabV3 模型。
在 Python 中获取示例输入图像的模型输出，并将其与 iOS 应用程序的输出进行比较。
构建一个新的 iOS 应用程序或重用一个 iOS 示例应用程序来加载转换后的模型。
将输入准备为模型期望的格式并处理模型输出。
完成 UI、重构、构建和运行应用程序以查看图像分割的实际效果。

先决条件

PyTorch 1.6 或 1.7
火炬视觉 0.7 或 0.8
Xcode 11 或 12

脚步

1.转换DeepLabV3模型用于iOS部署

在 iOS 上部署模型的第一步是将模型转换为TorchScript格式。

笔记

目前并非所有 PyTorch 模型都可以转换为 TorchScript，因为模型定义可能使用 TorchScript 中没有的语言功能，TorchScript 是 Python 的一个子集。有关更多详细信息，请参阅脚本和优化配方。

只需运行下面的脚本即可生成脚本模型deeplabv3_scripted.pt：

import torch

# use deeplabv3_resnet50 instead of deeplabv3_resnet101 to reduce the model size
model = torch.hub.load('pytorch/vision:v0.8.0', 'deeplabv3_resnet50', pretrained=True)
model.eval()

scriptedm = torch.jit.script(model)
torch.jit.save(scriptedm, "deeplabv3_scripted.pt")

生成的deeplabv3_scripted.pt模型文件的大小应该在 168MB 左右。理想情况下，在将模型部署到 iOS 应用程序之前，还应该对模型进行量化以显着减小尺寸并加快推理速度。要对量化有一个大致的了解，请参阅量化配方和那里的资源链接。我们将在以后的教程或秘籍中详细介绍如何将称为训练后静态量化的量化工作流正确应用于DeepLabV3 模型。

2.在Python中获取模型的示例输入和输出

现在我们有了一个脚本化的 PyTorch 模型，让我们用一些示例输入进行测试，以确保模型在 iOS 上正常工作。首先，让我们编写一个 Python 脚本，使用该模型进行推理并检查输入和输出。对于 DeepLabV3 模型的这个示例，我们可以重用步骤 1 和DeepLabV3 模型中心站点中的代码。将以下代码片段添加到上面的代码中：

from PIL import Image
from torchvision import transforms
input_image = Image.open("deeplab.jpg")
preprocess = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])

input_tensor = preprocess(input_image)
input_batch = input_tensor.unsqueeze(0)
with torch.no_grad():
    output = model(input_batch)['out'][0]

print(input_batch.shape)
print(output.shape)

从这里下载deeplab.jpg并运行上面的脚本以查看模型的输入和输出的形状：

torch.Size([1, 3, 400, 400])
torch.Size([21, 400, 400])

因此，如果您向iOS 上的模型提供大小为 400x400的相同图像输入deeplab.jpg，则模型的输出应具有 [21, 400, 400] 大小。您还应该至少打印出输入和输出的实际数据的开始部分，以便在下面的第 4 步中用于与模型在 iOS 应用程序中运行时的实际输入和输出进行比较。

3. 构建一个新的 iOS 应用程序或重用示例应用程序并加载模型

首先，按照iOS 模型准备教程的第 3 步，在启用 PyTorch Mobile 的 Xcode 项目中使用我们的模型。由于本教程中使用的 DeepLabV3 模型和 PyTorch HelloWorld iOS 示例中使用的 MobileNet v2 模型都是计算机视觉模型，因此您可以选择以HelloWorld 示例 repo作为模板来重用加载模型和处理数据的代码输入和输出。

现在让我们将步骤 2 中使用的deeplabv3_scripted.pt和deeplab.jpg添加到 Xcode 项目中，并将ViewController.swift修改为类似：

class ViewController: UIViewController {
    var image = UIImage(named: "deeplab.jpg")!

    override func viewDidLoad() {
        super.viewDidLoad()
    }

    private lazy var module: TorchModule = {
        if let filePath = Bundle.main.path(forResource: "deeplabv3_scripted",
              ofType: "pt"),
            let module = TorchModule(fileAtPath: filePath) {
            return module
        } else {
            fatalError("Can't load the model file!")
        }
    }()
}

然后在行返回模块处设置断点并构建并运行应用程序。应用程序应该在断点处停止，这意味着步骤 1 中的脚本模型已成功加载到 iOS 上。

4. 处理模型输入和输出以进行模型推理

在上一步中加载模型后，让我们验证它是否适用于预期的输入并可以生成预期的输出。由于 DeepLabV3 模型的模型输入是图像，与 HelloWorld 示例中的 MobileNet v2 相同，我们将重用HelloWorld 中TorchModule.mm文件中的部分代码进行输入处理。更换predictImage的方法实现TorchModule.mm用下面的代码：

- (unsigned char*)predictImage:(void*)imageBuffer {
    // 1. the example deeplab.jpg size is size 400x400 and there are 21 semantic classes
    const int WIDTH = 400;
    const int HEIGHT = 400;
    const int CLASSNUM = 21;

    at::Tensor tensor = torch::from_blob(imageBuffer, {1, 3, WIDTH, HEIGHT}, at::kFloat);
    torch::autograd::AutoGradMode guard(false);
    at::AutoNonVariableTypeMode non_var_type_mode(true);

    // 2. convert the input tensor to an NSMutableArray for debugging
    float* floatInput = tensor.data_ptr<float>();
    if (!floatInput) {
        return nil;
    }
    NSMutableArray* inputs = [[NSMutableArray alloc] init];
    for (int i = 0; i < 3 * WIDTH * HEIGHT; i++) {
        [inputs addObject:@(floatInput[i])];
    }

    // 3. the output of the model is a dictionary of string and tensor, as
    // specified at https://pytorch.org/hub/pytorch_vision_deeplabv3_resnet101
    auto outputDict = _impl.forward({tensor}).toGenericDict();

    // 4. convert the output to another NSMutableArray for easy debugging
    auto outputTensor = outputDict.at("out").toTensor();
    float* floatBuffer = outputTensor.data_ptr<float>();
    if (!floatBuffer) {
      return nil;
    }
    NSMutableArray* results = [[NSMutableArray alloc] init];
    for (int i = 0; i < CLASSNUM * WIDTH * HEIGHT; i++) {
      [results addObject:@(floatBuffer[i])];
    }

    return nil;
}

笔记

模型输出是 DeepLabV3 模型的字典，因此我们使用toGenericDict来正确提取结果。对于其他模型，模型输出也可能是单个张量或张量元组等。

通过上面显示的代码更改，您可以在填充输入和结果的两个 for 循环之后设置断点，并将它们与您在步骤 2 中看到的模型输入和输出数据进行比较，以查看它们是否匹配。对于在 iOS 和 Python 上运行的模型的相同输入，您应该获得相同的输出。

到目前为止，我们所做的只是确认我们感兴趣的模型可以像在 Python 中一样在我们的 iOS 应用程序中编写脚本并正确运行。到目前为止，我们在 iOS 应用程序中使用模型的步骤消耗了大部分（如果不是大部分）应用程序开发时间，类似于数据预处理是典型机器学习项目中最繁重的提升。

5. 完成 UI、重构、构建和运行应用程序

现在我们已准备好完成应用程序和 UI，以将处理后的结果作为新图像实际查看。输出处理代码应该是这样的，添加到TorchModule.mm中Step 4代码片段的最后——记得先去掉return nil这一行；暂时放在那里以使代码构建和运行：

// see the 20 semantic classes link in Introduction
const int DOG = 12;
const int PERSON = 15;
const int SHEEP = 17;

NSMutableData* data = [NSMutableData dataWithLength:
    sizeof(unsigned char) * 3 * WIDTH * HEIGHT];
unsigned char* buffer = (unsigned char*)[data mutableBytes];
// go through each element in the output of size [WIDTH, HEIGHT] and
// set different color for different classnum
for (int j = 0; j < WIDTH; j++) {
    for (int k = 0; k < HEIGHT; k++) {
        // maxi: the index of the 21 CLASSNUM with the max probability
        int maxi = 0, maxj = 0, maxk = 0;
        float maxnum = -100000.0;
        for (int i = 0; i < CLASSNUM; i++) {
            if ([results[i * (WIDTH * HEIGHT) + j * WIDTH + k] floatValue] > maxnum) {
                maxnum = [results[i * (WIDTH * HEIGHT) + j * WIDTH + k] floatValue];
                maxi = i; maxj = j; maxk = k;
            }
        }
        int n = 3 * (maxj * width + maxk);
        // color coding for person (red), dog (green), sheep (blue)
        // black color for background and other classes
        buffer[n] = 0; buffer[n+1] = 0; buffer[n+2] = 0;
        if (maxi == PERSON) buffer[n] = 255;
        else if (maxi == DOG) buffer[n+1] = 255;
        else if (maxi == SHEEP) buffer[n+2] = 255;
    }
}
return buffer;

这里的实现基于对DeepLabV3模型的理解，该模型为宽*高的输入图像输出大小为[21, width, height]的张量。width*height 输出数组中的每个元素都是一个介于 0 到 20 之间的值（对于介绍中描述的总共 21 个语义标签），该值用于设置特定颜色。这里分割的颜色编码是基于概率最高的类，你可以扩展你自己数据集中所有类的颜色编码。

输出处理后，还需要调用一个辅助函数来转换RGB缓冲到一个UIImage的实例上显示的UIImageView。您可以参考代码仓库中UIImageHelper.mm中定义的示例代码convertRGBBufferToUIImage。

此应用的 UI 也与 HelloWorld 的 UI 类似，只是您不需要UITextView来显示图像分类结果。您还可以添加两个按钮Segment和Restart如代码库中所示运行模型推理并在显示分割结果后显示原始图像。

我们可以运行应用程序之前的最后一步是将所有部分连接在一起。修改ViewController.swift文件以使用在repo中重构并更改为segmentImage 的 predictImage和您构建的辅助函数，如ViewController.swift中的 repo 中的示例代码所示。将按钮连接到动作，你应该很高兴。