微信小程序中使用TensorFlowJS从环境搭建到模型训练及推理模型得到预测结果

1、小程序端环境准备

app.json

  "plugins": {"tfjsPlugin": {"version": "0.2.0","provider": "wx6afed118d9e81df9"}}

package.json

  "dependencies": {"@tensorflow-models/posenet": "^2.2.2","@tensorflow/tfjs-backend-webgl": "3.5.0","@tensorflow/tfjs-converter": "3.5.0","@tensorflow/tfjs-core": "^3.5.0","@tensorflow/tfjs-layers": "^4.22.0","fetch-wechat": "^0.0.3","lottie-miniprogram": "^1.0.12"}

终端执行

Microsoft Windows [版本 10.0.19045.6093]
(c) Microsoft Corporation。保留所有权利。E:\AAASelfProjectGit\myWxProject> npm i

在微信开发者工具中点击工具->构建npm

2、训练模型

python环境

python          3.8.20
protobuf        3.20.3
numpy           1.22.0
tensorflowjs    3.7.0
tensorflow      2.13.0

训练代码 （使用手写数字数据集，keras自带minist）

import tensorflow as tfmnist = tf.keras.datasets.mnist(x_train, y_train),(x_test, y_test) = mnist.load_data()x_train, x_test = x_train / 255.0, x_test / 255.0model = tf.keras.models.Sequential([tf.keras.layers.Flatten(input_shape=(28, 28)),tf.keras.layers.Dense(128, activation='relu'),tf.keras.layers.Dropout(0.2),  tf.keras.layers.Dense(10, activation='softmax')])model.compile(optimizer='adam',loss='sparse_categorical_crossentropy',metrics=['accuracy'])model.fit(x_train, y_train, epochs=5)model.evaluate(x_test, y_test)model.save('D:\\mnist.h5')

训练完成后.h5文件会在D盘，下面代码将.h5转换成tensorflowjs所需要的.json格式

import os
import subprocessh5_model_path = "D:\\mnist.h5"
output_dir = "D:\\"
os.makedirs(output_dir, exist_ok=True)# 使用绝对路径调用（需替换为你的实际路径）
command = ["python","D:\\anaconda\\envs\\ckm\\Scripts\\tensorflowjs_converter.exe",  # Windows 路径示例"--input_format=keras",h5_model_path,output_dir
]try:subprocess.run(command, check=True)print("转换成功！")
except subprocess.CalledProcessError as e:print(f"转换失败，错误代码: {e.returncode}")

转换成功后会得到两个文件

将两个文件上传到服务器，通过地址访问.json

3、小程序端代码预测

在js中引入并使用

var fetchWechat = require('fetch-wechat');
var tf = require('@tensorflow/tfjs-core');
var tfl = require('@tensorflow/tfjs-layers');
var webgl = require('@tensorflow/tfjs-backend-webgl');
var plugin = requirePlugin('tfjsPlugin');Page({async onReady() {//加载相机const camera = wx.createCameraContext(this)// 加载模型const net = await this.loadModel()this.setData({result: 'Loading'})let count = 0//每隔10帧获取一张相机捕捉到的图片const listener = camera.onCameraFrame((frame) => {count++if (count === 10) {if (net) {//对图片内容进行预测this.predict(net, frame)}count = 0}})listener.start()},//加载模型async loadModel() {const net = await tfl.loadLayersModel('https://你的服务器域名.com/model.json')net.summary()return net},async predict(net, frame) {try {const x = tf.tidy(() => {const imgTensor = tf.tensor3d(new Uint8Array(frame.data), [frame.height, frame.width, 4])const d = Math.floor((frame.height - frame.width) / 2)const imgSlice = tf.slice(imgTensor, [d, 0, 0], [frame.width, frame.width, 3])const imgResize = tf.image.resizeBilinear(imgSlice, [28, 28])return tf.mean(imgResize, 2) // [28, 28]})// 添加批次维度 [1, 28, 28]const input = tf.reshape(x, [1, ...x.shape])// 预测并处理结果const prediction = await net.predict(input)// 使用tf.topk替代argMaxconst {values, indices} = tf.topk(prediction, 1)const res = indices.dataSync()[0]this.setData({result: res})// 释放内存tf.dispose([x, input, prediction, values, indices])} catch (error) {console.error('预测错误:', error)this.setData({result: 'Error: ' + error.message})}}})

在wxml中展示{{result}}即可看到预测结果

<view class="landscape-container"><!-- 相机层（横屏适配） --><camera device-position="back" resolution="high" frame-size="large" style="width: 100%; height: 100vh;z-index:10;" catch:tap="cameraClick" id="myCamera"></camera><view style="position: absolute;bottom: 100px;z-index: 99;left: 50%;transform: translateX(-50%);font-size: 20px;font-weight: 800;color: white;">预测结果：{{result}}</view>
</view>