PyTorch 分割模型的轉換與 OpenCV 的啟動

目標

在本教程中，您將學習如何

• 轉換 PyTorch 分割模型
• 使用 OpenCV 執行轉換後的 PyTorch 模型
• 獲取 PyTorch 和 OpenCV DNN 模型的評估

我們將透過 FCN ResNet-50 架構的示例來探討以上列出的要點。

介紹

使用 OpenCV API 進行 PyTorch 分類 和分割模型轉換流程中涉及的關鍵點是相同的。第一步是使用 PyTorch torch.onnx.export 內建函式將模型轉換為 ONNX 格式。然後，將獲得的 .onnx 模型傳遞給 cv.dnn.readNetFromONNX，它會返回 cv.dnn.Net 物件，該物件已準備好用於 DNN 操作。

實踐

在本部分中，我們將介紹以下幾點

1. 建立分割模型轉換流程並提供推理
2. 評估和測試分割模型

如果您只想執行評估或測試模型流程，則可以跳過“模型轉換流程”部分。

模型轉換流程

本小節中的程式碼位於 dnn_model_runner 模組中，可以使用以下行執行

python -m dnn_model_runner.dnn_conversion.pytorch.segmentation.py_to_py_fcnresnet50

以下程式碼包含下面列出的步驟的描述

1. 例項化 PyTorch 模型
2. 將 PyTorch 模型轉換為 .onnx
3. 使用 OpenCV API 讀取轉換後的網路
4. 準備輸入資料
5. 提供推理
6. 從預測中獲取彩色掩碼
7. 視覺化結果

# 初始化 PyTorch FCN ResNet-50 模型
original_model = models.segmentation.fcn_resnet50(pretrained=True)
 
# 獲取轉換為 ONNX 的 PyTorch 模型的路徑
full_model_path = get_pytorch_onnx_model(original_model)
 
# 使用 OpenCV API 讀取轉換後的 .onnx 模型
opencv_net = cv2.dnn.readNetFromONNX(full_model_path)
print("OpenCV 模型已成功讀取。圖層 ID：\n", opencv_net.getLayerNames())
 
# 獲取預處理後的影像
img, input_img = get_processed_imgs("test_data/sem_segm/2007_000033.jpg")
 
# 獲取 OpenCV DNN 預測
opencv_prediction = get_opencv_dnn_prediction(opencv_net, input_img)
 
# 獲取原始 PyTorch ResNet50 預測
pytorch_prediction = get_pytorch_dnn_prediction(original_model, input_img)
 
pascal_voc_classes, pascal_voc_colors = read_colors_info("test_data/sem_segm/pascal-classes.txt")
 
# 獲取彩色分割掩碼
opencv_colored_mask = get_colored_mask(img.shape, opencv_prediction, pascal_voc_colors)
pytorch_colored_mask = get_colored_mask(img.shape, pytorch_prediction, pascal_voc_colors)
 
# 獲取 PASCAL VOC 顏色的調色盤
color_legend = get_legend(pascal_voc_classes, pascal_voc_colors)
 
cv2.imshow('PyTorch 彩色掩碼', pytorch_colored_mask)
cv2.imshow('OpenCV DNN 彩色掩碼', opencv_colored_mask)
cv2.imshow('顏色圖例', color_legend)
 
cv2.waitKey(0)

為了提供模型推理，我們將使用來自 PASCAL VOC 驗證資料集的下圖

目標分割結果是

對於 PASCAL VOC 顏色解碼及其與預測掩碼的對映，我們還需要 pascal-classes.txt 檔案，其中包含 PASCAL VOC 類和相應顏色的完整列表。

讓我們透過預訓練的 PyTorch FCN ResNet-50 的示例更深入地瞭解每個程式碼步驟

• 例項化 PyTorch FCN ResNet-50 模型

# 初始化 PyTorch FCN ResNet-50 模型
original_model = models.segmentation.fcn_resnet50(pretrained=True)

• 將 PyTorch 模型轉換為 ONNX 格式

# 定義用於進一步轉換的模型儲存的目錄
onnx_model_path = "models"
# 定義進一步轉換的模型的名稱
onnx_model_name = "fcnresnet50.onnx"
 
# 建立用於進一步轉換的模型的目錄
os.makedirs(onnx_model_path, exist_ok=True)
 
# 獲取轉換為模型的完整路徑
full_model_path = os.path.join(onnx_model_path, onnx_model_name)
 
# 生成模型輸入以構建圖
generated_input = Variable(
torch.randn(1, 3, 500, 500)
)
 
# 模型匯出為 ONNX 格式
torch.onnx.export(
original_model,
generated_input,
full_model_path,
verbose=True,
input_names=["input"],
output_names=["output"],
opset_version=11
)

此步驟的程式碼與分類轉換的情況沒有區別。因此，在成功執行上述程式碼後，我們將獲得 models/fcnresnet50.onnx。

• 讀取轉換後的網路，使用 cv.dnn.readNetFromONNX 將上一步中獲得的 ONNX 模型傳遞給它

# 使用 OpenCV API 讀取轉換後的 .onnx 模型
opencv_net = cv2.dnn.readNetFromONNX(full_model_path)

• 準備輸入資料

# 讀取影像
input_img = cv2.imread(img_path, cv2.IMREAD_COLOR)
input_img = input_img.astype(np.float32)
 
# 目標影像大小
img_height = input_img.shape[0]
img_width = input_img.shape[1]
 
# 定義預處理引數
mean = np.array([0.485, 0.456, 0.406]) * 255.0
scale = 1 / 255.0
std = [0.229, 0.224, 0.225]
 
# 準備輸入 blob 以適應模型輸入
# 1. 減去均值
# 2. 縮放以將畫素值設定為 0 到 1
input_blob = cv2.dnn.blobFromImage(
image=input_img,
scalefactor=scale,
size=(img_width, img_height), # 影像目標大小
mean=mean,
swapRB=True, # BGR -> RGB
crop=False # 中心裁剪
)
# 3. 除以 std
input_blob[0] /= np.asarray(std, dtype=np.float32).reshape(3, 1, 1)

在此步驟中，我們讀取影像並使用 cv2.dnn.blobFromImage 函式準備模型輸入，該函式返回 4 維 blob。應該注意的是，首先在 cv2.dnn.blobFromImage 中減去均值，然後才縮放畫素值。因此，mean 乘以 255.0 以重現原始影像預處理順序

img /= 255.0
img -= [0.485, 0.456, 0.406]
img /= [0.229, 0.224, 0.225]

• OpenCV cv.dnn_Net 推理

# 設定 OpenCV DNN 輸入
opencv_net.setInput(preproc_img)
 
# OpenCV DNN 推理
out = opencv_net.forward()
print("OpenCV DNN 分割預測：\n")
print("* 形狀：", out.shape)
 
# 獲取預測類的 ID
out_predictions = np.argmax(out[0], axis=0)

在上述程式碼執行後，我們將獲得以下輸出

OpenCV DNN 分割預測
* 形狀：(1, 21, 500, 500)

21 個預測通道中的每個通道（其中 21 表示 PASCAL VOC 類的數量）都包含機率，這些機率指示畫素對應於 PASCAL VOC 類的可能性。

• PyTorch FCN ResNet-50 模型推理

original_net.eval()
preproc_img = torch.FloatTensor(preproc_img)
 
with torch.no_grad()
    # 獲取每個類的未歸一化機率
out = original_net(preproc_img)['out']
 
print("\nPyTorch 分割模型預測：\n")
print("* 形狀：", out.shape)
 
# 獲取預測類的 ID
out_predictions = out[0].argmax(dim=0)

在啟動上述程式碼後，我們將獲得以下輸出

PyTorch 分割模型預測
* 形狀：torch.Size([1, 21, 366, 500])

PyTorch 預測還包含對應於每個類預測的機率。

• 從預測中獲取彩色掩碼

# 將掩碼值轉換為 PASCAL VOC 顏色
processed_mask = np.stack([colors[color_id] for color_id in segm_mask.flatten()])
 
# 將掩碼重塑為 3 通道影像
processed_mask = processed_mask.reshape(mask_height, mask_width, 3)
processed_mask = cv2.resize(processed_mask, (img_width, img_height), interpolation=cv2.INTER_NEAREST).astype(
np.uint8)
 
# 將彩色掩碼從 BGR 轉換為 RGB，以便與 PASCAL VOC 顏色相容
processed_mask = cv2.cvtColor(processed_mask, cv2.COLOR_BGR2RGB)

在此步驟中，我們將分割掩碼中的機率對映到預測類的適當顏色。讓我們看一下結果

對於模型的擴充套件評估，我們可以使用 dnn_model_runner 模組的 py_to_py_segm 指令碼。此模組部分將在下一小節中描述。

模型評估

dnn/samples 中提出的 dnn_model_runner 模組允許在 PASCAL VOC 資料集上執行完整的評估流程，並測試以下 PyTorch 分割模型的執行情況

• FCN ResNet-50
• FCN ResNet-101

此列表也可以透過進一步的適當評估流程配置來擴充套件。

評估模式

以下行表示在評估模式下執行模組

python -m dnn_model_runner.dnn_conversion.pytorch.segmentation.py_to_py_segm --model_name <pytorch_segm_model_name>

從列表中選擇的分割模型將被讀取到 OpenCV cv.dnn_Net 物件中。PyTorch 和 OpenCV 模型的評估結果（畫素精度、平均 IoU、推理時間）將被寫入日誌檔案。推理時間值也將在圖表中描述，以概括獲得的模型資訊。

必要的評估配置在 test_config.py 中定義

@dataclass
class TestSegmConfig
frame_size: int = 500
img_root_dir: str = "./VOC2012"
img_dir: str = os.path.join(img_root_dir, "JPEGImages/")
img_segm_gt_dir: str = os.path.join(img_root_dir, "SegmentationClass/")
    # 減少 val：https://github.com/shelhamer/fcn.berkeleyvision.org/blob/master/data/pascal/seg11valid.txt
segm_val_file: str = os.path.join(img_root_dir, "ImageSets/Segmentation/seg11valid.txt")
colour_file_cls: str = os.path.join(img_root_dir, "ImageSets/Segmentation/pascal-classes.txt")

這些值可以根據選擇的模型流程進行修改。

要啟動 PyTorch FCN ResNet-50 的評估，請執行以下行

python -m dnn_model_runner.dnn_conversion.pytorch.segmentation.py_to_py_segm --model_name fcnresnet50

測試模式

以下行表示在測試模式下執行模組，該模式提供模型推理的步驟

python -m dnn_model_runner.dnn_conversion.pytorch.segmentation.py_to_py_segm --model_name <pytorch_segm_model_name> --test True --default_img_preprocess <True/False> --evaluate False

這裡的 default_img_preprocess 鍵定義您是否希望使用某些特定值來引數化模型測試過程，或者使用預設值，例如 scale、mean 或 std。

測試配置在 test_config.py TestSegmModuleConfig 類中表示

@dataclass
class TestSegmModuleConfig
segm_test_data_dir: str = "test_data/sem_segm"
test_module_name: str = "segmentation"
test_module_path: str = "segmentation.py"
input_img: str = os.path.join(segm_test_data_dir, "2007_000033.jpg")
model: str = ""
 
frame_height: str = str(TestSegmConfig.frame_size)
frame_width: str = str(TestSegmConfig.frame_size)
scale: float = 1.0
mean: List[float] = field(default_factory=lambda: [0.0, 0.0, 0.0])
std: List[float] = field(default_factory=list)
crop: bool = False
rgb: bool = True
classes: str = os.path.join(segm_test_data_dir, "pascal-classes.txt")

預設影像預處理選項在 default_preprocess_config.py 中定義

pytorch_segm_input_blob = {
    "mean": ["123.675", "116.28", "103.53"],
    "scale": str(1 / 255.0),
    "std": ["0.229", "0.224", "0.225"],
    "crop": "False",
    "rgb": "True"
}

模型測試的基礎在 samples/dnn/segmentation.py 中表示。segmentation.py 可以使用 --input 中提供的轉換模型和 cv2.dnn.blobFromImage 的填充引數自主執行。

要從頭開始重現“模型轉換流程”中描述的 OpenCV 步驟與 dnn_model_runner，請執行以下行

python -m dnn_model_runner.dnn_conversion.pytorch.segmentation.py_to_py_segm --model_name fcnresnet50 --test True --default_img_preprocess True --evaluate False