IOS 上的圖像分割 DEEPLABV3
介紹
語(yǔ)義圖像分割是一項(xiàng)計(jì)算機(jī)視覺任務(wù)，它使用語(yǔ)義標(biāo)簽來(lái)標(biāo)記輸入圖像的特定區(qū)域。PyTorch 語(yǔ)義圖像分割DeepLabV3 模型可用于標(biāo)記具有20 個(gè)語(yǔ)義類的圖像區(qū)域，例如，自行車、公共汽車、汽車、狗和人。圖像分割模型在自動(dòng)駕駛和場(chǎng)景理解等應(yīng)用中非常有用。
在本教程中，我們將提供有關(guān)如何在 iOS 上準(zhǔn)備和運(yùn)行 PyTorch DeepLabV3 模型的分步指南，帶您從開始擁有一個(gè)您可能想要在 iOS 上使用的模型到最終擁有一個(gè)完整的模型。使用該模型的 iOS 應(yīng)用程序。我們還將介紹有關(guān)如何檢查下一個(gè)最喜歡的預(yù)訓(xùn)練 PyTorch 模型是否可以在 iOS 上運(yùn)行以及如何避免陷阱的實(shí)用和一般技巧。
筆記
在學(xué)習(xí)本教程之前，您應(yīng)該查看適用于 iOS 的 PyTorch Mobile，并快速嘗試一下PyTorch iOS HelloWorld示例應(yīng)用程序。本教程將超越圖像分類模型，通常是部署在移動(dòng)設(shè)備上的第一種模型。本教程的完整代碼存儲(chǔ)庫(kù)可在此處獲得。
學(xué)習(xí)目標(biāo)
在本教程中，您將學(xué)習(xí)如何：
為 iOS 部署轉(zhuǎn)換 DeepLabV3 模型。
在 Python 中獲取示例輸入圖像的模型輸出，并將其與 iOS 應(yīng)用程序的輸出進(jìn)行比較。
構(gòu)建一個(gè)新的 iOS 應(yīng)用程序或重用一個(gè) iOS 示例應(yīng)用程序來(lái)加載轉(zhuǎn)換后的模型。
將輸入準(zhǔn)備為模型期望的格式并處理模型輸出。
完成 UI、重構(gòu)、構(gòu)建和運(yùn)行應(yīng)用程序以查看圖像分割的實(shí)際效果。
先決條件
PyTorch 1.6 或 1.7
火炬視覺 0.7 或 0.8
Xcode 11 或 12
腳步
1.轉(zhuǎn)換DeepLabV3模型用于iOS部署
在 iOS 上部署模型的第一步是將模型轉(zhuǎn)換為TorchScript格式。
筆記
目前并非所有 PyTorch 模型都可以轉(zhuǎn)換為 TorchScript，因?yàn)槟Ｐ投x可能使用 TorchScript 中沒有的語(yǔ)言功能，TorchScript 是 Python 的一個(gè)子集。有關(guān)更多詳細(xì)信息，請(qǐng)參閱腳本和優(yōu)化配方。
只需運(yùn)行下面的腳本即可生成腳本模型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模型文件的大小應(yīng)該在 168MB 左右。理想情況下，在將模型部署到 iOS 應(yīng)用程序之前，還應(yīng)該對(duì)模型進(jìn)行量化以顯著減小尺寸并加快推理速度。要對(duì)量化有一個(gè)大致的了解，請(qǐng)參閱量化配方和那里的資源鏈接。我們將在以后的教程或秘籍中詳細(xì)介紹如何將稱為訓(xùn)練后靜態(tài)量化的量化工作流正確應(yīng)用于DeepLabV3 模型。
2.在Python中獲取模型的示例輸入和輸出
現(xiàn)在我們有了一個(gè)腳本化的 PyTorch 模型，讓我們用一些示例輸入進(jìn)行測(cè)試，以確保模型在 iOS 上正常工作。首先，讓我們編寫一個(gè) Python 腳本，使用該模型進(jìn)行推理并檢查輸入和輸出。對(duì)于 DeepLabV3 模型的這個(gè)示例，我們可以重用步驟 1 和DeepLabV3 模型中心站點(diǎn)中的代碼。將以下代碼片段添加到上面的代碼中：
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并運(yùn)行上面的腳本以查看模型的輸入和輸出的形狀：
torch.Size([1, 3, 400, 400])
torch.Size([21, 400, 400])
因此，如果您向iOS 上的模型提供大小為 400x400的相同圖像輸入deeplab.jpg，則模型的輸出應(yīng)具有 [21, 400, 400] 大小。您還應(yīng)該至少打印出輸入和輸出的實(shí)際數(shù)據(jù)的開始部分，以便在下面的第 4 步中用于與模型在 iOS 應(yīng)用程序中運(yùn)行時(shí)的實(shí)際輸入和輸出進(jìn)行比較。
3. 構(gòu)建一個(gè)新的 iOS 應(yīng)用程序或重用示例應(yīng)用程序并加載模型
首先，按照iOS 模型準(zhǔn)備教程的第 3 步，在啟用 PyTorch Mobile 的 Xcode 項(xiàng)目中使用我們的模型。由于本教程中使用的 DeepLabV3 模型和 PyTorch HelloWorld iOS 示例中使用的 MobileNet v2 模型都是計(jì)算機(jī)視覺模型，因此您可以選擇以HelloWorld 示例 repo作為模板來(lái)重用加載模型和處理數(shù)據(jù)的代碼輸入和輸出。
現(xiàn)在讓我們將步驟 2 中使用的deeplabv3_scripted.pt和deeplab.jpg添加到 Xcode 項(xiàng)目中，并將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!")
}
}()
}
然后在行返回模塊處設(shè)置斷點(diǎn)并構(gòu)建并運(yùn)行應(yīng)用程序。應(yīng)用程序應(yīng)該在斷點(diǎn)處停止，這意味著步驟 1 中的腳本模型已成功加載到 iOS 上。
4. 處理模型輸入和輸出以進(jìn)行模型推理
在上一步中加載模型后，讓我們驗(yàn)證它是否適用于預(yù)期的輸入并可以生成預(yù)期的輸出。由于 DeepLabV3 模型的模型輸入是圖像，與 HelloWorld 示例中的 MobileNet v2 相同，我們將重用HelloWorld 中TorchModule.mm文件中的部分代碼進(jìn)行輸入處理。更換predictImage的方法實(shí)現(xiàn)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來(lái)正確提取結(jié)果。對(duì)于其他模型，模型輸出也可能是單個(gè)張量或張量元組等。
通過上面顯示的代碼更改，您可以在填充輸入和結(jié)果的兩個(gè) for 循環(huán)之后設(shè)置斷點(diǎn)，并將它們與您在步驟 2 中看到的模型輸入和輸出數(shù)據(jù)進(jìn)行比較，以查看它們是否匹配。對(duì)于在 iOS 和 Python 上運(yùn)行的模型的相同輸入，您應(yīng)該獲得相同的輸出。
到目前為止，我們所做的只是確認(rèn)我們感興趣的模型可以像在 Python 中一樣在我們的 iOS 應(yīng)用程序中編寫腳本并正確運(yùn)行。到目前為止，我們?cè)?iOS 應(yīng)用程序中使用模型的步驟消耗了大部分（如果不是大部分）應(yīng)用程序開發(fā)時(shí)間，類似于數(shù)據(jù)預(yù)處理是典型機(jī)器學(xué)習(xí)項(xiàng)目中最繁重的提升。
5. 完成 UI、重構(gòu)、構(gòu)建和運(yùn)行應(yīng)用程序
現(xiàn)在我們已準(zhǔn)備好完成應(yīng)用程序和 UI，以將處理后的結(jié)果作為新圖像實(shí)際查看。輸出處理代碼應(yīng)該是這樣的，添加到TorchModule.mm中Step 4代碼片段的最后——記得先去掉return nil這一行；暫時(shí)放在那里以使代碼構(gòu)建和運(yùn)行：
// 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;
這里的實(shí)現(xiàn)基于對(duì)DeepLabV3模型的理解，該模型為寬*高的輸入圖像輸出大小為[21, width, height]的張量。width*height 輸出數(shù)組中的每個(gè)元素都是一個(gè)介于 0 到 20 之間的值（對(duì)于介紹中描述的總共 21 個(gè)語(yǔ)義標(biāo)簽），該值用于設(shè)置特定顏色。這里分割的顏色編碼是基于概率最高的類，你可以擴(kuò)展你自己數(shù)據(jù)集中所有類的顏色編碼。
輸出處理后，還需要調(diào)用一個(gè)輔助函數(shù)來(lái)轉(zhuǎn)換RGB緩沖到一個(gè)UIImage的實(shí)例上顯示的UIImageView。您可以參考代碼倉(cāng)庫(kù)中UIImageHelper.mm中定義的示例代碼convertRGBBufferToUIImage。
此應(yīng)用的 UI 也與 HelloWorld 的 UI 類似，只是您不需要UITextView來(lái)顯示圖像分類結(jié)果。您還可以添加兩個(gè)按鈕Segment和Restart如代碼庫(kù)中所示運(yùn)行模型推理并在顯示分割結(jié)果后顯示原始圖像。
我們可以運(yùn)行應(yīng)用程序之前的最后一步是將所有部分連接在一起。修改ViewController.swift文件以使用在repo中重構(gòu)并更改為segmentImage 的 predictImage和您構(gòu)建的輔助函數(shù)，如ViewController.swift中的 repo 中的示例代碼所示。將按鈕連接到動(dòng)作，你應(yīng)該很高興。
現(xiàn)在，當(dāng)您在 iOS 模擬器或?qū)嶋H iOS 設(shè)備上運(yùn)行該應(yīng)用程序時(shí)，您將看到以下屏幕：
