看过美剧《硅谷》的同学都知道人工智能识别热狗曾是硅谷最赚钱的技术之一。去年 HBO 发布了官方的 Not Hotdog 应使用,支持 iOS 和 Android 平台,据说是使用 TensorFlow、Keras 和 React Native 打造的,但是源码没有公开。
最近入门了机器学习的少量边边角角,仿制了一个 Not Hotdog 应使用。代码主要用 React Native 和 TensorFlow Mobile ,训练模型用 Yolo v2 Tiny。
shaqian/Not-Hotdog
是时候表演真正的技术了 —— 画框
偷懒使用了 COCO 2014 数据集。COCO 一共有 80 种物体类别,其中一个就是热狗 (hot dog)。我写了个脚本导出所有的热狗图片(共 1200 多张)并生成 yolo 格式的标注:
shaqian/Not-Hotdog/blob/master/yolo/coco2yolo.py
用的是 darknet 的 这个 fork 。
darknet.exe partial yolov2-tiny.cfg yolov2-tiny.weights yolov2-tiny.conv.13 13
而后整理数据集,主要是生成 obj.data 和 obj.names 文件,可参考:
https://timebutt.github.io/static/how-to-train-yolov2-to-detect-custom-objects/
修改 .cfg 文件:
将本来的 80 个类别改成 1 个类别。.cfg 文件和初始 weights 请见 GitHub repo 。
最后训练模型:
darknet.exe detector train data\obj.data yolov2-tiny-hotdog.cfg yolov2-tiny.conv.13
我使用的 darkflow 转换:
flow --model ../yolov2-tiny-hotdog.cfg --load ../yolov2-tiny-hotdog_final.weights --savepb
转换后的 .pb 约有 44MB。使用 TensorFlow repo里面的 quantization 脚本 压缩一下,体积可以减小到 11MB。
python3 tensorflow/tools/quantization/quantize_graph.py --input=yolov2-tiny-hotdog.pb --output_node_names=output --output=quantized_yolov2-tiny-hotdog.pb --mode=weights
出于方便,调使用 TensorFlow Mobile 接口是自己实现的 native 板块。
iOS 端的实现代码在 ./react-native-NotHotdog/ios/NotHotdog/TensorflowManager.mm
Android 端的实现代码在 ./react-native-NotHotdog/android/app/src/main/java/com/nothotdog/tensorflowmanager
参考了以下 repo 特此感谢:
TensorFlow Android example
TensorFlow iOS example
yolov2_tf_ios
cd react-native-NotHotdog/iospod install
cd react-native-NotHotdog/npm install
react-native run-ios
真实设施:
请参考 Running your app on iOS devices。
因为 react-native-camera 的少量问题,Android 端拍照要比 iOS 慢很多,有时间可以改进一下。
GitHub repo 里的 hotdog.apk 可直接安装。
模拟器:
react-native run-android
真实设施:
请参考 Running your app on Android devices。