AI Pocket —— 能跑深度模型的小程序

简要介绍

「AI Pocket」是一款能跑深度模型的微信小程序,无需服务器后端,利用 TensorFlow.js(以下简称 tfjs)将精简模型加载到手机上,并实现模型预测功能,甚至可以在手机上进行模型的训练。

附上小程序二维码,欢迎大家扫码体验。代码开源在 GitHub 上并且同步更新,仓库地址为 HunterXuan/wx-tfjs-demo,欢迎有兴趣的同学一起加功能 & 改 BUG。

AI Pocket 小程序码

踩的坑

小程序目前更多是图像识别相关的模型,用到了小程序的 onCameraFrame API 获取帧数据,然后转换成模型输入进行预测。然而,由于小程序的 onCameraFrame 获取到的帧数据与所展示的不一致,而且是在不同的设备上(甚至相同设备的前后摄像头)对原始帧数据的处理方式都不一样,要想得到准确的预测结果,真叫人头大。

好在目前摸索出来一定的裁剪套路,也经过简单的测试:IOS 无论前后摄像头,都从上往下裁;安卓前置摄像头从上往下裁,后置摄像头从下往上裁。这是用手上仅有的安卓和 IOS 设备测试的结果,可能无法覆盖所有设备,如果遇到预测结果非常离谱的情况,欢迎到 GitHub 上提 Issues。

开发故事

导火索

某一天,大佬和我聊到 TensorFlow.js,问我有没可能搬到小程序上,说是没准可以做个爆款。经过一番魔改,我成功把 tfjs 移植到了微信小程序上。

然而,由于小程序 API 诸多的不完善,以及没有办法优雅地用上 webgl,模型预测速度实在太慢,就暂时搁置了。我简单总结了我的魔改过程,发表了一篇文章 tfjs 移植到微信小程序,然后也就没有了爆款的下文。

首次开源 Demo

过了几个月,有人看到我写的一篇关于魔改 tfjs 的文章,留言表示很有兴趣,希望能开源。后面还有几个大学生因为毕设选了相关的题目,陆续来咨询我实现细节。

然后,我就整理了下旧代码并开源在了 GitHub 上,也就是目前 AI Pocket 的开源代码仓 HunterXuan/wx-tfjs-demo。同时,还发表了一篇博文 TensorFlowJS 移植微信小程序再次尝试,希望能给有需要的同学一些帮助。

改良优化

又过了几个月,同事给我分享了一篇文章,里面有提到 TensorFlow 官方对小程序的支持插件。我掐指一算,感觉可以继续搞搞。

我连夜瞄了 tfjs 最近的更新和小程序插件的代码,发现 tfjs 开始考虑对多平台的支持,而且小程序也开放了更多有利的 API。

于是,我就把之前开源的 Demo 再完善了下,用上了 webgl 加速,预测速度飞起,便有了现在的 AI Pocket。

把 AI 装进口袋

我希望把更多的 AI 模型装到 AI Pocket 上,在小程序上就可以跑机器学习/深度学习模型,打造一个「流弊」的 AI 口袋。