本文介紹了如何使用 TensorFlow 在智能機(jī)上（包括安卓和 iOS 設(shè)備）執(zhí)行實(shí)時(shí)單人姿態(tài)估計(jì)。

GitHub 地址：https://github.com/edvardHua/PoseEstimationForMobile

該 repo 使用 TensorFlow 實(shí)現(xiàn) CPM 和 Hourglass 模型。這里未使用常規(guī)的卷積，而是在模型內(nèi)部使用了反向卷積（又叫 Mobilenet V2），以便執(zhí)行實(shí)時(shí)推斷。

注：你可以修改網(wǎng)絡(luò)架構(gòu)，來(lái)訓(xùn)練更高 PCKh 的模型。架構(gòu)地址：https://github.com/edvardHua/PoseEstimationForMobile/tree/master/training/src

該庫(kù)包含：

• 訓(xùn)練 CPM 和 Hourglass 模型的代碼；
• 安卓 demo 的源代碼；
• IOS demo 的源代碼。

下面的 gif 是在 Mi Mix2s 上截取的（~60 FPS）

你可以下載以下 apk，在自己的設(shè)備上進(jìn)行測(cè)試。

• PoseEstimation-Mace.apk：https://raw./edvardHua/PoseEstimationForMobile/master/release/PoseEstimation-Mace.apk
• PoseEstimation-TFlite.apk：https://raw./edvardHua/PoseEstimationForMobile/master/release/PoseEstimation-TFlite.apk

訓(xùn)練

依賴項(xiàng)

• Python3
• TensorFlow >= 1.4
• Mace

數(shù)據(jù)集

訓(xùn)練數(shù)據(jù)集：https://drive.google.com/open?id=1zahjQWhuKIYWRRI2ZlHzn65Ug_jIiC4l

將其解壓縮，獲取以下文件結(jié)構(gòu)：

該訓(xùn)練數(shù)據(jù)集僅包含單人圖像，數(shù)據(jù)來(lái)源是 AI Challenger 競(jìng)賽。共包含 22446 個(gè)訓(xùn)練樣本和 1500 個(gè)測(cè)試樣本。

該 repo 作者使用 tf-pose-estimation 庫(kù)中的數(shù)據(jù)增強(qiáng)代碼將標(biāo)注遷移為 COCO 格式。tf-pose-estimation 庫(kù)：https://github.com/ildoonet/tf-pose-estimation

超參數(shù)

訓(xùn)練步驟中，使用 experiments 文件夾中的 cfg 文件傳輸超參數(shù)。

以下是 mv2_cpm.cfg 文件的內(nèi)容：

[Train]model: 'mv2_cpm'checkpoint: Falsedatapath: '/root/hdd/ai_challenger'imgpath: '/root/hdd/'visible_devices: '0, 1, 2'multiprocessing_num: 8max_epoch: 1000lr: '0.001'batchsize: 5decay_rate: 0.95input_width: 192input_height: 192n_kpoints: 14scale: 2modelpath: '/root/hdd/trained/mv2_cpm/models'logpath: '/root/hdd/trained/mv2_cpm/log'num_train_samples: 20000per_update_tensorboard_step: 500per_saved_model_step: 2000pred_image_on_tensorboard: True

該 cfg 文件覆蓋模型的所有參數(shù)，在 network_mv2_cpm.py 中仍有一些參數(shù)。

使用 nvidia-docker 訓(xùn)練

通過(guò)以下命令構(gòu)建 docker：

或者

docker pull edvardhua/single-pose

然后運(yùn)行以下命令，訓(xùn)練模型：

此外，它還在 port 6006 上創(chuàng)建了 tensorboard。確保安裝了 nvidia-docker。

按一般方法訓(xùn)練

1. 安裝依賴項(xiàng)：

cd trainingpip3 install -r requirements.txt

還需要安裝 cocoapi (https://github.com/cocodataset/cocoapi)。

2. 編輯 experiments 文件夾中的參數(shù)文件，它包含幾乎所有超參數(shù)和訓(xùn)練中需要定義的其他配置。之后，傳輸參數(shù)文件，開(kāi)始訓(xùn)練：

在 3 張英偉達(dá) 1080Ti 顯卡上經(jīng)過(guò) 12 個(gè)小時(shí)的訓(xùn)練后，該模型幾乎收斂。以下是對(duì)應(yīng)的 tensorboard 圖。

基準(zhǔn)（PCKh）

運(yùn)行以下命令，評(píng)估 PCKh 值。

python3 src/benchmark.py --frozen_pb_path=hourglass/model-360000.pb \--anno_json_path=/root/hdd/ai_challenger/ai_challenger_valid.json \--img_path=/root/hdd \--output_node_name=hourglass_out_3

預(yù)訓(xùn)練模型

• CPM：https://github.com/edvardHua/PoseEstimationForMobile/tree/master/release/cpm_model
• Hourglass：https://github.com/edvardHua/PoseEstimationForMobile/tree/master/release/hourglass_model

安卓 demo

由于 mace 框架，你可以使用 GPU 在安卓智能機(jī)上運(yùn)行該模型。

按照以下命令將模型轉(zhuǎn)換為 mace 格式：

然后根據(jù) mace 文檔的說(shuō)明，將模型集成到安卓設(shè)備中。

至于如何調(diào)用模型、解析輸出，可以參見(jiàn)安卓源代碼：https://github.com/edvardHua/PoseEstimationForMobile/tree/master/android_demo。

一些芯片的平均推斷時(shí)間基準(zhǔn)如下所示：

以下是該 repo 作者構(gòu)建該 demo 的環(huán)境：

• 操作系統(tǒng)：macOS 10.13.6（mace 目前不支持 windows）
• Android Studio：3.0.1
• NDK 版本：r16

在構(gòu)建 mace-demo 時(shí)，不同環(huán)境可能會(huì)遇到不同的錯(cuò)誤。為避免這種情況，作者建議使用 docker。

docker pull registry.cn-hangzhou.aliyuncs.com/xiaomimace/mace-dev-litedocker run -it --privileged -d --name mace-dev  --net=host  -v to/you/path/PoseEstimationForMobile/android_demo/demo_mace:/demo_mace  registry.cn-hangzhou.aliyuncs.com/xiaomimace/mace-dev-litedocker run -it --privileged -d --name mace-dev --net=host \ -v to/you/path/PoseEstimationForMobile/android_demo/demo_mace:/demo_mace \ registry.cn-hangzhou.aliyuncs.com/xiaomimace/mace-dev-lite# Enter to dockerdocker exec -it mace-dev bash# Exec command inside the dockercd /demo_mace && ./gradlew build

或者將模型轉(zhuǎn)換為 tflite：

然后，將 tflite 文件放在 android_demo/app/src/main/assets 中，修改 ImageClassifierFloatInception.kt 中的參數(shù)。

............// parameters need to modify in ImageClassifierFloatInception.kt/** * Create ImageClassifierFloatInception instance * * @param imageSizeX Get the image size along the x axis. * @param imageSizeY Get the image size along the y axis. * @param outputW The output width of model * @param outputH The output height of model * @param modelPath Get the name of the model file stored in Assets. * @param numBytesPerChannel Get the number of bytes that is used to store a single * color channel value. */ fun create( activity: Activity, imageSizeX: Int = 192, imageSizeY: Int = 192, outputW: Int = 96, outputH: Int = 96, modelPath: String = 'mv2-cpm.tflite', numBytesPerChannel: Int = 4 ): ImageClassifierFloatInception = ImageClassifierFloatInception( activity, imageSizeX, imageSizeY, outputW, outputH, modelPath, numBytesPerChannel)............

最后，將該項(xiàng)目導(dǎo)入 Android Studio，在智能機(jī)設(shè)備上運(yùn)行。

iOS Demo

首先，將模型轉(zhuǎn)換為 CoreML 模型：

然后，按照 PoseEstimation-CoreML 中的說(shuō)明來(lái)操作（https://github.com/tucan9389/PoseEstimation-CoreML）。