TensorFlow.js实现商标识别

在VsCode中利用TensorFlow.js结合迁移学习实现商标识别。

一、加载商标数据并可视化

数据保存在data文件夹下面，需要先在data文件夹下创建一个静态服务器，用于加载图片。

http-server data --cors

Available on:
  http://192.168.4.167:8080
  http://127.0.0.1:8080
Hit CTRL-C to stop the server

 编写获取图片的脚本文件。

const IMAGE_SIZE = 224;

const loadImg = (src) => {
    return new Promise(resolve => {
        const img = new Image();
        img.crossOrigin = "anonymous";
        img.src = src;
        img.width = IMAGE_SIZE;
        img.height = IMAGE_SIZE;
        img.onload = () => resolve(img);
    });
};
export const getInputs = async () => {
    const loadImgs = [];
    const labels = [];
    for (let i = 0; i < 30; i += 1) {
        ['android', 'apple', 'windows'].forEach(label => {
            const src = `http://127.0.0.1:8080/brand/train/${label}-${i}.jpg`;
            const img = loadImg(src);
            loadImgs.push(img);
            labels.push([
                label === 'android' ? 1 : 0,
                label === 'apple' ? 1 : 0,
                label === 'windows' ? 1 : 0,
            ]);
        });
    }
    const inputs = await Promise.all(loadImgs);
    return {
        inputs,
        labels,
    };
}

创建index.html文件，作为程序的入口文件，在index.html中利用script标签跳转到script.js文件，在script.js中编写主要代码。

加载图片数据。

import {getInputs} from "./data";

window.onload = async() =>{
    const {inputs, labels} = await getInputs();
    console.log(inputs, labels)

};

利用TensorFlow.js中的tfvis进行可视化。

    import * as tfvis from "@tensorflow/tfjs-vis"

    // 可视化图片
    const surface = tfvis.visor().surface({ name: '输入示例', styles: { height: 250 } });
    inputs.forEach(img => {
        surface.drawArea.appendChild(img);
    });

每行显示两个，旁边的滚动体可以拉动查看更多图片。

二、定义模型结构

加载MobileNet模型并截断所有的卷积池化操作，生成截断模型。并定义新的全连接层。

    import * as tf from "@tensorflow/tfjs"

    // mobilenet模型存放位置
    const MOBILENET_MODEL_PATH = 'http://127.0.0.1:8080/mobilenet/web_model/model.json';
    
    // 加载MobileNet模型
    const mobilenet = await tf.loadLayersModel(MOBILENET_MODEL_PATH);
    // 查看模型结构
    mobilenet.summary();
    // 截断mobilenet卷积操作
    const layer = mobilenet.getLayer('conv_pw_13_relu');
    const truncatedMobilenet = tf.model({
        inputs: mobilenet.inputs,
        outputs: layer.output
    });

    // 定义全连接层
    const model = tf.sequential();
    model.add(tf.layers.flatten({
        inputShape: layer.outputShape.slice(1)
    }));
    model.add(tf.layers.dense({
        units: 10,
        activation: 'relu'
    }));
    // 定义输出层
    model.add(tf.layers.dense({
        units: NUM_CLASSES,
        activation: 'softmax'
    }));
    // 配置损失函数和优化器
    model.compile({ loss: 'categoricalCrossentropy', optimizer: tf.train.adam() });

三、迁移学习下的模型训练

首先先定义一个工具类utils.js，用于处理输入到截断模型(mobilenet)中的数据。

import * as tf from '@tensorflow/tfjs';
// img格式转成tensor
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;
    });
}
// 图片文件转成img格式
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);
        };
    });
}

训练数据经过截断模型输出，转为可以用于自定义的全连接层的输入数据。

    // 先经过截断模型
    const { xs, ys } = tf.tidy(() => {
        const xs = tf.concat(inputs.map(imgEl =>truncatedMobilenet.predict(img2x(imgEl))));
        const ys = tf.tensor(labels);
        return { xs, ys };
    });
    // 截断模型的输出当成自定义模型的输入
    await model.fit(xs, ys, {
        epochs: 20,
        callbacks: tfvis.show.fitCallbacks(
            { name: '训练效果' },
            ['loss'],
            { callbacks: ['onEpochEnd'] }
        )
    });

可以看出训练损失值降得非常低，因为采用迁移模型，卷积层的参数是使用别人训练好的，这部分参数的训练结果是非常优秀的。

四、预测 

编写前端页面用于上传带预测图片，就是编写一个上传按钮。

<script src="script.js"></script>

<input type="file" onchange="predict(this.files[0])">

将预测图片先经过mobileNet预测，吐出来的结果再经过自定义模型预测。

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

        const index = pred.argMax(1).dataSync()[0];
        setTimeout(() => {
            alert(`预测结果：${BRAND_CLASSES[index]}`);
        }, 0);
    };

 

五、完整代码

index.html

<script src="script.js"></script>

<input type="file" onchange="predict(this.files[0])">

script.js 

import * as tf from "@tensorflow/tfjs"
import * as tfvis from "@tensorflow/tfjs-vis"
import {getInputs} from "./data";
import {img2x, file2img} from "./utils"

// mobilenet模型存放位置
const MOBILENET_MODEL_PATH = 'http://127.0.0.1:8080/mobilenet/web_model/model.json';

const NUM_CLASSES = 3;
const BRAND_CLASSES = ['android', 'apple', 'windows'];


window.onload = async() =>{
    // 加载图片
    const {inputs, labels} = await getInputs();
    // 可视化图片
    const surface = tfvis.visor().surface({ name: '输入示例', styles: { height: 250 } });
    inputs.forEach(img => {
        surface.drawArea.appendChild(img);
    });

    // 加载MobileNet模型
    const mobilenet = await tf.loadLayersModel(MOBILENET_MODEL_PATH);
    // 查看模型结构
    mobilenet.summary();
    // 截断mobilenet卷积操作
    const layer = mobilenet.getLayer('conv_pw_13_relu');
    const truncatedMobilenet = tf.model({
        inputs: mobilenet.inputs,
        outputs: layer.output
    });

    // 定义全连接层
    const model = tf.sequential();
    model.add(tf.layers.flatten({
        inputShape: layer.outputShape.slice(1)
    }));
    model.add(tf.layers.dense({
        units: 10,
        activation: 'relu'
    }));
    // 定义输出层
    model.add(tf.layers.dense({
        units: NUM_CLASSES,
        activation: 'softmax'
    }));
    // 配置损失函数和优化器
    model.compile({ loss: 'categoricalCrossentropy', optimizer: tf.train.adam() });


    // 先经过截断模型
    const { xs, ys } = tf.tidy(() => {
        const xs = tf.concat(inputs.map(imgEl => truncatedMobilenet.predict(img2x(imgEl))));
        const ys = tf.tensor(labels);
        return { xs, ys };
    });
    // 截断模型的输出当成自定义模型的输入
    await model.fit(xs, ys, {
        epochs: 20,
        callbacks: tfvis.show.fitCallbacks(
            { name: '训练效果' },
            ['loss'],
            { callbacks: ['onEpochEnd'] }
        )
    });

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

        const index = pred.argMax(1).dataSync()[0];
        setTimeout(() => {
            alert(`预测结果：${BRAND_CLASSES[index]}`);
        }, 0);
    };


};

data.js 

const IMAGE_SIZE = 224;

const loadImg = (src) => {
    return new Promise(resolve => {
        const img = new Image();
        img.crossOrigin = "anonymous";
        img.src = src;
        img.width = IMAGE_SIZE;
        img.height = IMAGE_SIZE;
        img.onload = () => resolve(img);
    });
};
export const getInputs = async () => {
    const loadImgs = [];
    const labels = [];
    for (let i = 0; i < 30; i += 1) {
        ['android', 'apple', 'windows'].forEach(label => {
            const src = `http://127.0.0.1:8080/brand/train/${label}-${i}.jpg`;
            const img = loadImg(src);
            loadImgs.push(img);
            labels.push([
                label === 'android' ? 1 : 0,
                label === 'apple' ? 1 : 0,
                label === 'windows' ? 1 : 0,
            ]);
        });
    }
    const inputs = await Promise.all(loadImgs);
    return {
        inputs,
        labels,
    };
}

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);
        };
    });
}