快速注册

用 Tensorflow.js 做了一个动漫分类的功能（二）

ZERO开发 2024-05-08 14:09:35 广东

前言：

前面已经通过采集拿到了图片，并且也手动对图片做了标注。接下来就要通过 Tensorflow.js 基于 mobileNet 训练模型，最后就可以实现在采集中对图片进行自动分类了。

这种功能在应用场景里就比较多了，比如图标素材站点，用户通过上传一个图标，系统会自动匹配出相似的图标，还有二手平台，用户通过上传闲置物品图片，平台自动给出分类等，这些也都是前期对海量图片进行了标注训练而得到一个损失率极低的模型。下面就通过简答的代码实现一个小的动漫分类。

环境：

Node

Http-Server

Parcel

Tensorflow

编码：

1. 训练模型

1.1. 创建项目，安装依赖包

npm install @tensorflow/tfjs --legacy-peer-deps
npm install @tensorflow/tfjs-node-gpu --legacy-peer-deps

1.2. 全局安装 Http-Server

npm install i http-server

1.3. 下载 mobileNet 模型文件 (网上有下载)

1.4. 根目录下启动 Http 服务 (开启跨域)，用于 mobileNet 和训练结果的模型可访问

http-server --cors -p 8080

1.5. 创建训练执行脚本 run.js

const tf = require('@tensorflow/tfjs-node-gpu');

const getData = require('./data');
const TRAIN_PATH = './动漫分类/train';
const OUT_PUT = 'output';
const MOBILENET_URL = 'http://127.0.0.1:8080/data/mobilenet/web_model/model.json';

(async () => {
  const { ds, classes } = await getData(TRAIN_PATH, OUT_PUT);
  console.log(ds, classes);
  //引入别人训练好的模型
  const mobilenet = await tf.loadLayersModel(MOBILENET_URL);
  //查看模型结构
  mobilenet.summary();

  const model = tf.sequential();
  //截断模型，复用了86个层
  for (let i = 0; i < 86; ++i) {
    const layer = mobilenet.layers[i];
    layer.trainable = false;
    model.add(layer);
  }
  //降维，摊平数据
  model.add(tf.layers.flatten());
  //设置全连接层
  model.add(tf.layers.dense({
    units: 10,
    activation: 'relu'//设置激活函数，用于处理非线性问题
  }));

  model.add(tf.layers.dense({
    units: classes.length,
    activation: 'softmax'//用于多分类问题
  }));
  //设置损失函数，优化器
  model.compile({
    loss: 'sparseCategoricalCrossentropy',
    optimizer: tf.train.adam(),
    metrics:['acc']
  });

  //训练模型
  await model.fitDataset(ds, { epochs: 20 });
  //保存模型
  await model.save(`file://${process.cwd()}/${OUT_PUT}`);
})();

1.6. 创建图片与 Tensor 转换库 data.js

const fs = require('fs');
const tf = require("@tensorflow/tfjs-node-gpu");

const img2x = (imgPath) => {
  const buffer = fs.readFileSync(imgPath);
  //清除数据
  return tf.tidy(() => {
    //把图片转成tensor
    const imgt = tf.node.decodeImage(new Uint8Array(buffer), 3);
    //调整图片大小
    const imgResize = tf.image.resizeBilinear(imgt, [224, 224]);
    //归一化
    return imgResize.toFloat().sub(255 / 2).div(255 / 2).reshape([1, 224, 224, 3]);
  });
}

const getData = async (traindir, output) => {
  let classes = fs.readdirSync(traindir, 'utf-8');
  fs.writeFileSync(`./${output}/classes.json`, JSON.stringify(classes));
  const data = [];
  classes.forEach((dir, dirIndex) => {
    fs.readdirSync(`${traindir}/${dir}`)
      .filter(n => n.match(/jpg$/))
      .slice(0, 1000)
      .forEach(filename => {
        const imgPath = `${traindir}/${dir}/${filename}`;

        data.push({ imgPath, dirIndex });
      });
  });

  console.log(data);

  //打乱训练顺序，提高准确度
  tf.util.shuffle(data);

  const ds = tf.data.generator(function* () {
    const count = data.length;
    const batchSize = 32;
    for (let start = 0; start < count; start += batchSize) {
      const end = Math.min(start + batchSize, count);
      console.log('当前批次', start);
      yield tf.tidy(() => {
        const inputs = [];
        const labels = [];
        for (let j = start; j < end; ++j) {
          const { imgPath, dirIndex } = data[j];
          const x = img2x(imgPath);
          inputs.push(x);
          labels.push(dirIndex);
        }
        const xs = tf.concat(inputs);
        const ys = tf.tensor(labels);
        return { xs, ys };
      });
    }
  });

  return { ds, classes };
}

module.exports = getData;

1.7. 运行执行文件

node run.js

2. 调用模型

2.1. 全局安装 parcel

npm install i parcel

2.2. 创建页面 index.html

<script src="script.js"></script>
<input type="file" onchange="predict(this.files[0])">
<br>

2.3. 创建模型调用预测脚本 script.js

import * as tf from '@tensorflow/tfjs';
import { img2x, file2img } from './utils';

const MODEL_PATH = 'http://127.0.0.1:8080/t7';
const CLASSES = ["假面骑士","奥特曼","海贼王","火影忍者","龙珠"];


window.onload = async () => {
    const model = await tf.loadLayersModel(MODEL_PATH + '/output/model.json');

    window.predict = async (file) => {
        const img = await file2img(file);
        document.body.appendChild(img);
        const pred = tf.tidy(() => {
            const x = img2x(img);
            return model.predict(x);
        });

        const index = pred.argMax(1).dataSync()[0];
        console.log(pred.argMax(1).dataSync());

        let predictStr = "";
        if (typeof CLASSES[index] == 'undefined') {
            predictStr = BRAND_CLASSES[index];
        } else {
            predictStr = CLASSES[index];
        }

        setTimeout(() => {
            alert(`预测结果：${predictStr}`);
        }, 0);
    };
};

2.4. 创建图片 tensor 格式转换库 utils.js

import * as tf from '@tensorflow/tfjs';

export function img2x(imgEl){
    return tf.tidy(() => {
        const input = tf.browser.fromPixels(imgEl)
            .toFloat()
            .sub(255 / 2)
            .div(255 / 2)
            .reshape([1, 224, 224, 3]);
        return input;
    });
}

export function file2img(f) {
    return new Promise(resolve => {
        const reader = new FileReader();
        reader.readAsDataURL(f);
        reader.onload = (e) => {
            const img = document.createElement('img');
            img.src = e.target.result;
            img.width = 224;
            img.height = 224;
            img.onload = () => resolve(img);
        };
    });
}

2.5. 打包项目并运行