FFmpeg中基於深度學習模型的圖像處理filter:dnn_processing介紹(2)

dnn_processing是FFmpeg中的一個視頻filter,可以支持所有基於深度學習模型的圖像處理算法,即輸入和輸出都是AVFrame,而處理過程使用的是深度學習模型。爲什麼要開發這樣一個filter,因爲作爲FFmpeg DNN模塊的maintainer,希望可以有更多的人來使用這個模塊,提出意見和批評,乃至發出patch來改善這個模塊,dnn_processing就是一個很好的使用者入手功能。至於支持視頻分析功能的filter,則還在計劃中,不過我不準備太着急開始編碼,一來還要考慮如何支持異步建立流水線,如何啓用batch size,從而最大化的用好系統的並行計算能力(感謝和pengfei的討論);二來則是希望再等待DNN模塊一段時間,以更好的夯實基礎,完善DNN接口,完善後端支持。

在上一篇點擊進入介紹瞭如何使用dnn_processing來完成針對灰度圖的sobel算子的調用,其輸入輸出的格式是grayf32。本篇將介紹dnn_processing如何完成sr(超分辨率)和derain(去除雨點)這兩個filter的功能,從而演示對yuv和rgb格式的支持。

  • 重新編譯FFmpeg

FFmpeg對深度學習模型的支持有兩個後端,一是調用TensorFlow C動態庫來完成,二是調用內部C代碼來實現。前者被稱爲tensorflow後端(將簡寫爲tf後端),後者被稱爲native後端。要啓用tf後端,必須在編譯FFmpeg的機器上有tensorflow c開發庫(包括.h文件和.so文件),而且在configure的時候需要傳入相應的參數。而任何情況下的編譯,native後端總是被自動啓用。這兩個後端本文都會涉及,所以,我們必須重新編譯FFmpeg。更多細節,請見上一篇介紹。

  • 支持sr功能

首先,我們要準備sr需要的深度學習模型,sr filter支持兩個模型,分別是srcnn和espcn,具體區別將在後續使用時候介紹。tf模型文件可以通過https://github.com/HighVoltageRocknRoll/srhttps://github.com/XueweiMeng/sr中的腳本生成,得到tf模型後再用FFmpeg目錄樹中的腳本生成native模型文件,具體如下。

# 演示目的,我們放到/tmp目錄下。
# 系統重啓後,/tmp目錄下的所有內容都會消失。
$ cd /tmp/

$ git clone https://github.com/HighVoltageRocknRoll/sr.git
$ cd sr/

# 下面這條命令會在當前目錄下生成tf模型,文件名是srcnn.pb
$ python generate_header_and_model.py --model=srcnn --ckpt_path=checkpoints/srcnn/
$ ls srcnn.pb -s
36 srcnn.pb

# 下面這條命令會在當前目錄下生成tf模型,文件名是espcn.pb
$ python generate_header_and_model.py --model=espcn --ckpt_path=checkpoints/espcn 
$ ls espcn.pb -s
88 espcn.pb

# 然後去ffmpeg目錄樹下
$ cd /path_to_your_ffmpeg_source_tree/
$ cd tools/python/

# 下面命令生成native模型文件srcnn.model
$ python convert.py /tmp/sr/srcnn.pb
$ ls -s srcnn.model 
36 srcnn.model

# 下面命令生成native模型文件espcn.model
$ python convert.py /tmp/sr/espcn.pb
$ ls -s espcn.model 
84 espcn.model

srcnn模型是第一個基於深度學習的超分辨率實現,其原理是首先將一個低分辨率圖片放大,然後,針對放大後的高分辨率圖片的Y通道基於深度學習做處理,UV通道保持不變,處理前後的分辨率保持不變,最後得到高質量的高分辨率圖片。具體命令如下所示。

$ cd /path_to_ffmpeg_build_path/

# 下面命令的輸入是一張480p的圖片,分辨率爲720x480。
# 命令行中使用了三個filter,分別是format,scale和dnn_processing,
# format用來確保到達dnn_processing的格式是YUV格式,
# scale則是將低分辨率圖片的長和寬分別放大2倍,
# dnn_processing中的model參數指出模型文件,dnn_backend指出使用tf後端,
# 而input和output則是模型的輸入輸出變量名字,在用python腳本搭建模型的時候指定。
# 最後,命令行的輸出圖片文件名是srcnn.jpg,其大小是1440*960。
$ ./ffmpeg -i 480p.jpg -vf \
format=yuv420p, \
scale=w=iw*2:h=ih*2, \
dnn_processing=dnn_backend=tensorflow:model=/tmp/sr/srcnn.pb:input=x:output=y \
srcnn.jpg

# 修改dnn_processing中的dnn_backend和model參數,
# 就可以切換到native後端執行,其他都保持不變。
# 記得一定要用剛纔生成的srcnn.model文件。
$ ./ffmpeg -i 480p.jpg -vf \
format=yuv420p, \
scale=w=iw*2:h=ih*2, \
dnn_processing=dnn_backend=native:model=/your_ffmpeg_source_tree/tools/python/srcnn.model:input=x:output=y \
srcnn.jpg 

espcn模型則比較簡潔,不需要提前進行放大,其模型的輸入是低分辨率圖片,輸出就是高分辨率圖片,也是隻針對Y通道的處理,UV通道調用swscale模塊(參數SWS_BICUBIC)進行放大。

$ cd /path_to_ffmpeg_build_path/

# 這裏就不需要用到scale了,其他參數解釋見前面介紹。
$ ./ffmpeg -i 480p.jpg -vf \
format=yuv420p, \
dnn_processing=dnn_backend=tensorflow:model=/tmp/sr/espcn.pb:input=x:output=y \
espcn.jpg

# 下面命令行用native後端執行。
$ ./ffmpeg -i 480p.jpg -vf \
format=yuv420p, \
dnn_processing=dnn_backend=native:model=/your_ffmpeg_source_tree/tools/python/espcn.model:input=x:output=y \
espcn.jpg
#
  • 支持derain功能

derain的深度學習模型的獲取方法如下所示,最後得到 tf模型文件can.pb 和 native模型文件can.model。

$ cd /tmp/
$ git clone https://github.com/XueweiMeng/derain_filter
$ cd derain_filter/scripts/
$ ./export_model.sh
$ ls -s ../can.pb  
148 ../can.pb

$ cd /path_to_your_ffmpeg_source_tree/
$ cd tools/python/
$ python convert.py /tmp/derain_filter/can.pb 
$ ls -s can.model 
108 can.model

derain模型支持的格式是rgb24,所以,FFmpeg命令行如下所示。

# 輸入rain.jpg是帶雨點的圖片,
# 用到了兩個filter,format和dnn_processing,
# 具體參數已經在前面有介紹。
$ ./ffmpeg -i rain.jpg -vf \
format=rgb24, \
dnn_processing=dnn_backend=tensorflow:model=can.pb:input=x:output=y \
derain.jpg

# 下面命令使用native後端執行
$ ./ffmpeg -i rain.jpg -vf \
format=rgb24, \
dnn_processing=dnn_backend=native:model=/your_ffmpeg_source_tree/tools/python/can.model:input=x:output=y \
derain.jpg

以上介紹完畢。如果以後出現處理YUV全部通道的深度學習模型成爲主流的情況,將會增加相應的支持。

以上內容是本人業餘時間興趣之作,限於水平,差錯難免,僅代表個人觀點,和本人任職公司無關。

本文首發於微信公衆號:那遁去的一

發表評論
所有評論
還沒有人評論,想成為第一個評論的人麼? 請在上方評論欄輸入並且點擊發布.
相關文章