OpenCV DNN 模塊-風格遷移

本文主要介紹OpenCV的DNN模塊的使用。OpenCV的DNN模塊自從contrib倉庫開始，就是只支持推理，不支持訓練。但是僅僅只是推理方面，也夠強大了。現在OpenCV已經支持TensorFlow、Pytorch/Torch、Caffe、DarkNet等模型的讀取。本文們就以風格遷移為例，來看一下OpenCV DNN模塊的用法。

相比于復雜而耗時的模型訓練過程，模型推理就顯得簡單多了。簡單來說，過程就是：

1. 加載模型

2. 輸入圖像預處理（跟訓練過程一樣的方式，增強除外）

3. 模型推理

1. 加載模型

因為OpenCV只支持推理，所以首先你需要有一個訓練好的模型。OpenCV支持所有主流框架的大部分模型。從OpenCV的readNet系列函數就可以看出來：

• readNetFromCaffe

• readNetFromTensorflow

• readNetFromTorch

• readNetFromDarknet

• readNetFromONNX

• readNetFromModelOptimizer

本文所用風格遷移模型是李飛飛的文章<<Perceptual Losses for Real-Time Style Transfer and Super-Resolution>>開源的Torch/Lua的模型，地址在這里：https://github.com/jcjohnson/fast-neural-style。他們提供了十種風格遷移的模型，模型的下載腳本在：https://github.com/jcjohnson/fast-neural-style/blob/master/models/download_style_transfer_models.sh。顯然這里需要用OpenCV的readNetFromTorch函數去加載模型，由于模型較多，這里提供的函數可以選擇加載指定的模型：

import cv2

model_base_dir = "/cvpy/style_transfer/models/"
d_model_map = {
    1: "udnie",
    2: "la_muse",
    3: "the_scream",
    4: "candy",
    5: "mosaic",
    6: "feathers",
    7: "starry_night"
}

def get_model_from_style(style: int):
    """
    加載指定風格的模型
    :param style: 模型編碼
    :return: model
    """
    model_name = d_model_map.get(style, "mosaic")
    model_path = model_base_dir + model_name + ".t7"
    model = cv2.dnn.readNetFromTorch(model_path)
    return model

2. 圖像預處理

在OpenCV中，輸入給模型的圖像需要首先被構建成一個4維的Blob，看到Blob這個詞感覺是受到了Caffe的影響。在構建Blob的時候會做一些諸如resize、歸一化和縮放之類的簡單預處理。OpenCV提供的函數為:

blobFromImage(image, scalefactor=None, size=None, mean=None, swapRB=None, crop=None, ddepth=None)

這個函數在構建Blob的之前會先做如下計算：

(image - mean) * scalefactor。

函數中的swapRB參數的含義是swap Blue and Red channels，干的是cvtColor(image, cv2.COLOR_BGR2RGB)的事情。

本文的風格遷移所需要做的圖像預處理很簡單，只是三通道分別減去均值即可。代碼如下：

(h, w) = img.shape[:2]
blob = cv2.dnn.blobFromImage(img, 1.0, (w, h), (103.939, 116.779, 123.680), swapRB=False, crop=False)

3. 模型推理

模型推理過程就是神經網絡模型進行一次前向傳播，在OpenCV中，用以下可讀性非常強的兩行代碼即可完成：

net.setInput(blob)
output = net.forward()

把第一節構建的blob輸入給模型，然后執行一次前向傳播。

得到輸出output再做一些處理使得我們可以更好的可視化圖像：

# reshape輸出結果, 將減去的平均值加回來，并交換各顏色通道
output = output.reshape((3, output.shape[2], output.shape[3]))
output[0] += 103.939
output[1] += 116.779
output[2] += 123.680
output = output.transpose(1, 2, 0)

效果展示

找一張測試圖片，選擇不同的風格，試一下效果。

想用自己的圖片風格遷移一下嗎？cvpy.net網站剛部署成功，來試試吧。