基于本地TensorFlow Lite和AI云服務(wù)的文檔矯正功能

什么是文檔圖片彎曲矯正？

用戶有時(shí)候會(huì)有這種需求痛點(diǎn)：看書時(shí)特別喜歡其中的某些段落，想摘抄下來，但發(fā)現(xiàn)篇幅很長(zhǎng)，用手機(jī)拍攝的話書本又不能完全鋪平。

手持一份裝訂過的合同，著急發(fā)給老板，卻沒有工具拆開裝訂釘。拍出來的圖里頁(yè)面是彎曲的，里邊的文字看起來很費(fèi)勁...

你也不可能隨時(shí)隨地帶著掃描儀。那該怎么辦呢？市面上大部分 APP 是把視角歪掉的圖片矯正為正視圖，如下圖：

在很多場(chǎng)合確實(shí)可以幫助到用戶。但是它只是解決了透視變化的問題。如果是彎曲的頁(yè)面，還是不能很好的矯正。

WPS 對(duì)這個(gè)需求很上心，希望可以讓用戶做到隨手一拍就獲得完美鋪平的電子文檔。歷經(jīng)幾個(gè)月的開發(fā)周期，結(jié)合了本地 AI 和云端 AI 能力打造了一個(gè)功能：文檔矯正。這個(gè)貼心的功能能讓剛才那些棘手的問題輕松被稿定。直接看效果吧。還是剛才那張拍攝的書本照片，經(jīng)過文檔彎曲矯正后。就可以立刻得到鋪平的圖片，甚至連頁(yè)面區(qū)域都給你自動(dòng)裁剪好了，背景干擾全無，完美！

這個(gè)彎曲矯正對(duì)于 OCR 識(shí)別也有很大的幫助，因?yàn)楹芏?OCR 識(shí)別對(duì)于文本彎曲的情況下識(shí)別會(huì)有很多亂碼，經(jīng)過彎曲矯正變成規(guī)規(guī)整整的文本行后，OCR 識(shí)別準(zhǔn)確度會(huì)大大提高。

以上就是我們新推出的這個(gè)功能的應(yīng)用場(chǎng)景和使用效果。那么下面我們就來聊一聊為什么這個(gè)功能背后是本地 AI 和云端 AI 的結(jié)合。

云+端 AI ，共建文檔矯正功能

目前其他第三方以及我們 WPS 里上一個(gè)版本的文檔矯正都是通過檢測(cè)出文檔的邊緣或者四個(gè)角，然后施行透視變化進(jìn)行矯正的。其中的計(jì)算量雖然也會(huì)比較大，但現(xiàn)在的移動(dòng)端計(jì)算能力還能夠覆蓋，實(shí)時(shí)檢測(cè)預(yù)覽起來也還比較流暢。然而彎曲矯正背后的計(jì)算量非常大，我們不得不將其部署在計(jì)算集群上，通過網(wǎng)絡(luò)傳輸壓縮和加密的原圖以及相應(yīng)的參數(shù)到服務(wù)器上進(jìn)行計(jì)算，然后將計(jì)算結(jié)果傳輸回用戶的手機(jī)端再呈現(xiàn)。這個(gè)過程原本也可以設(shè)計(jì)成全部在云端計(jì)算，那樣我們的工作量就會(huì)少很多，但是為了獲得更好的用戶體驗(yàn)，能夠讓用戶直觀的看到實(shí)時(shí)檢測(cè)的文檔區(qū)域，我們將文檔檢測(cè)的部分放在了本地。在用戶開啟相機(jī)進(jìn)行預(yù)覽時(shí)我們的深度學(xué)習(xí)模型就在后臺(tái)進(jìn)行運(yùn)行，逐幀去檢測(cè)文檔區(qū)域，并將檢測(cè)結(jié)果渲染在屏幕上。為了實(shí)現(xiàn)準(zhǔn)確穩(wěn)定的文檔區(qū)域檢測(cè)，我們將傳統(tǒng)的邊緣檢測(cè)算法升級(jí)為深度學(xué)習(xí)模型，隨著模型的不斷迭代，雖然識(shí)別效果越來越好，但模型體積也越來越大，計(jì)算量也越來越大。為了兼顧中低端手機(jī)的運(yùn)算能力，我們?cè)诰W(wǎng)絡(luò)模型設(shè)計(jì)時(shí)做了相應(yīng)的優(yōu)化，參考 SqueezeNet 的結(jié)構(gòu)設(shè)計(jì)了一個(gè) 7 層的卷積神經(jīng)網(wǎng)絡(luò)替代最初采用的更深的開源預(yù)訓(xùn)練網(wǎng)絡(luò)，雖然網(wǎng)絡(luò)變淺導(dǎo)致準(zhǔn)確度上稍微有些損失，但速度有大幅提升，內(nèi)存占用也減少不少，這使得低端手機(jī)也能達(dá)到每秒數(shù)幀的速度，而高端手機(jī)則可以達(dá)到每秒 30 幀滿幀的速度。同時(shí)為了彌補(bǔ)準(zhǔn)確度損失帶來的檢測(cè)不準(zhǔn)而引起的抖動(dòng)，我們?cè)诰W(wǎng)絡(luò)外邊增加了濾波器，讓檢測(cè)結(jié)果在視覺上看起來更穩(wěn)定。

整套網(wǎng)絡(luò)和代碼完成后，我們發(fā)現(xiàn)其實(shí) SDK 包的體積已經(jīng)有一些大了，除了對(duì)代碼和依賴庫(kù)的精簡(jiǎn)，我們還需要對(duì)推理框架進(jìn)行精簡(jiǎn)。我們需要一套速度快而又輕量級(jí)的移動(dòng)端推理框架。我們當(dāng)時(shí)做了很多選擇和嘗試，最后選擇了 Tensorflow Lite。Tensorflow Lite 是 Google 出品的對(duì)移動(dòng)端非常友好的深度學(xué)習(xí)框架，其架構(gòu)設(shè)計(jì)和訓(xùn)練部署思路都和 Tensorflow 非常相像，同時(shí)又非常小巧，容易上手。我們直接在 PC 架構(gòu)的訓(xùn)練機(jī)器上訓(xùn)練好模型，經(jīng)過 Google 提供的工具 TOCO 進(jìn)行轉(zhuǎn)化，即可獲得一個(gè)體積很小而且速度很快的手機(jī)端模型，而運(yùn)行的效果和準(zhǔn)確度卻沒有打多少折扣。從模型訓(xùn)練到移動(dòng)端部署的流程非常通暢，這讓我們非常喜歡 Tensorflow Lite 這個(gè)框架，畢竟我們可以減少工程部署的時(shí)間，而把精力更多放在算法優(yōu)化和模型優(yōu)化上。

除了獲取文檔區(qū)域，為了實(shí)現(xiàn)彎曲矯正，我們還設(shè)計(jì)了一整套復(fù)雜的算法系統(tǒng)，經(jīng)過評(píng)估，我們確定很難在移動(dòng)端進(jìn)行那樣高強(qiáng)度的計(jì)算，于是我們選擇了計(jì)算能力更強(qiáng)的云服務(wù)方案。核心算法部署在計(jì)算集群上可以獲得更充沛的算力，同時(shí)也為該業(yè)務(wù)算法的優(yōu)化和升級(jí)預(yù)留了比較大的空間，畢竟現(xiàn)在只是第一版，以后我們還會(huì)對(duì)其進(jìn)行持續(xù)升級(jí)和優(yōu)化。

目前我們的整體設(shè)計(jì)框架是將本地相應(yīng)的 AI 能力封裝成一個(gè) SDK 提供給業(yè)務(wù)方的客戶端，客戶端通過 API 接口調(diào)用 SDK 里相應(yīng)的能力，數(shù)據(jù)和參數(shù)也是通過 API 接口和 SDK 進(jìn)行交互，同時(shí) SDK 也接管了訪問云服務(wù)器的能力，統(tǒng)一的進(jìn)行管理，在保障了信息安全的同時(shí)也有相應(yīng)的容錯(cuò)和異常處理能力。這套系統(tǒng)的流程圖如下所示：

云和端的選擇以及未來的期許

因?yàn)樾枨蟮奶厥庑院筒煌惴ǖ膹?fù)雜程度以及平臺(tái)的計(jì)算能力評(píng)估，我們因地適宜的選擇了讓一部分計(jì)算在本地利用 TensorFlow Lite 進(jìn)行，一部分計(jì)算放在云端計(jì)算集群上進(jìn)行。相信隨著手機(jī)算力不斷的提升，以后也許單純依靠端上的算力就能完成如此復(fù)雜程度的計(jì)算。那一天應(yīng)該很快到來。當(dāng)然，隨著 5G 的普及，也許以后云端的計(jì)算借助更高帶寬極快速的網(wǎng)絡(luò)傳輸能力也能夠體驗(yàn)到端上的體驗(yàn)。到時(shí)候普通用戶再也感受不到云和端的區(qū)別，那將會(huì)給更復(fù)雜更強(qiáng)大的算法更多落地的機(jī)會(huì)。我們很期待這些條件更快滿足，我們也會(huì)努力帶給大家更多強(qiáng)大好用的算法功能。