Springboot整合Java DL4J实现交通标志识别系统全过程

traffic_sign_dataset/
├── images/
│   ├── sign1.jpg
│   ├── sign2.jpg
│   ├──...
├── labels/
│   ├── sign1.txt
│   ├── sign2.txt
│   ├──...

<dependency>
    <groupId>org.deeplearning4j</groupId>
    <artifactId>deeplearning4j-core</artifactId>
    <version>1.0.0-beta7</version>
</dependency>
<dependency>
    <groupId>org.deeplearning4j</groupId>
    <artifactId>deeplearning4j-nn</artifactId>
    <version>1.0.0-beta7</version>
</dependency>
<dependency>
    <groupId>org.deeplearning4j</groupId>
    <artifactId>deeplearning4j-ui</artifactId>
    <version>1.0.0-beta7</version>
</dependency>
<dependency>
    <groupId>org.springframework.boot</groupId>
    <artifactId>spring-boot-starter-web</artifactId>
</dependency>

import org.datavec.image.loader.NativeImageLoader;
import org.deeplearning4j.datasets.iterator.impl.ListDataSetIterator;
import org.nd4j.linalg.api.ndarray.INDArray;
import org.nd4j.linalg.dataset.DataSet;
import org.nd4j.linalg.dataset.api.preprocessor.DataNormalization;
import org.nd4j.linalg.dataset.api.preprocessor.ImagePreProcessingScaler;

import java.io.File;
import java.util.ArrayList;
import java.util.List;

public class DataLoader {

    public static ListDataSetIterator loadData(String dataDirectory) {
        // 加载图像文件
        File imageDirectory = new File(dataDirectory + "/images");
        NativeImageLoader imageLoader = new NativeImageLoader(32, 32, 3);
        List<INDArray> images = new ArrayList<>();
        for (File imageFile : imageDirectory.listFiles()) {
            INDArray image = imageLoader.asMatrix(imageFile);
            images.add(image);
        }

        // 加载标签文件
        File labelDirectory = new File(dataDirectory + "/labels");
        List<Integer> labels = new ArrayList<>();
        for (File labelFile : labelDirectory.listFiles()) {
            // 假设标签文件中每行只有一个数字，表示标签类别
            int label = Integer.parseInt(FileUtils.readFileToString(labelFile));
            labels.add(label);
        }

        // 创建数据集
        DataSet dataSet = new DataSet(images.toArray(new INDArray[0]), labels.stream().mapToDouble(i -> i).toArray());

        // 数据归一化
        DataNormalization scaler = new ImagePreProcessingScaler(0, 1);
        scaler.fit(dataSet);
        scaler.transform(dataSet);

        return new ListDataSetIterator(dataSet, 32);
    }
}

import org.deeplearning4j.nn.conf.ConvolutionMode;
import org.deeplearning4j.nn.conf.NeuralNetConfiguration;
import org.deeplearning4j.nn.conf.layers.ConvolutionLayer;
import org.deeplearning4j.nn.conf.layers.DenseLayer;
import org.deeplearning4j.nn.conf.layers.OutputLayer;
import org.deeplearning4j.nn.multilayer.MultiLayerNetwork;
import org.deeplearning4j.nn.weights.WeightInit;
import org.nd4j.linalg.activations.Activation;
import org.nd4j.linalg.lossfunctions.LossFunctions;

public class TrafficSignRecognitionModel {

    public static MultiLayerNetwork buildModel() {
        NeuralNetConfiguration.Builder builder = new NeuralNetConfiguration.Builder()
               .seed(12345)
               .weightInit(WeightInit.XAVIER)
               .updater(org.deeplearning4j.nn.weights.WeightInit.XAVIER)
               .l2(0.0005)
               .list();

        // 添加卷积层
        builder.layer(0, new ConvolutionLayer.Builder(5, 5)
               .nIn(3)
               .stride(1, 1)
               .nOut(32)
               .activation(Activation.RELU)
               .convolutionMode(ConvolutionMode.Same)
               .build());

        // 添加池化层
        builder.layer(1, new org.deeplearning4j.nn.conf.layers.SubsamplingLayer.Builder(org.deeplearning4j.nn.conf.layers.PoolingType.MAX)
               .kernelSize(2, 2)
               .stride(2, 2)
               .build());

        // 添加更多卷积层和池化层
        builder.layer(2, new ConvolutionLayer.Builder(5, 5)
               .nOut(64)
               .activation(Activation.RELU)
               .convolutionMode(ConvolutionMode.Same)
               .build());
        builder.layer(3, new org.deeplearning4j.nn.conf.layers.SubsamplingLayer.Builder(org.deeplearning4j.nn.conf.layers.PoolingType.MAX)
               .kernelSize(2, 2)
               .stride(2, 2)
               .build());

        // 添加全连接层
        builder.layer(4, new DenseLayer.Builder()
               .nOut(1024)
               .activation(Activation.RELU)
               .build());

        // 添加输出层
        builder.layer(5, new OutputLayer.Builder(LossFunctions.LossFunction.NEGATIVELOGLIKELIHOOD)
               .nOut(10) // 假设共有 10 种交通标志类别
               .activation(Activation.SOFTMAX)
               .build());

        return new MultiLayerNetwork(builder.build());
    }

    public static void trainModel(MultiLayerNetwork model, ListDataSetIterator iterator) {
        model.init();
        for (int epoch = 0; epoch < 10; epoch++) {
            model.fit(iterator);
            iterator.reset();
        }
    }
}

import org.deeplearning4j.nn.multilayer.MultiLayerNetwork;
import org.nd4j.linalg.api.ndarray.INDArray;
import org.nd4j.linalg.dataset.api.preprocessor.DataNormalization;
import org.nd4j.linalg.dataset.api.preprocessor.ImagePreProcessingScaler;
import org.nd4j.linalg.factory.Nd4j;

import java.io.File;

public class Prediction {

    public static int predict(MultiLayerNetwork model, File imageFile) {
        // 加载图像并进行预处理
        NativeImageLoader imageLoader = new NativeImageLoader(32, 32, 3);
        INDArray image = imageLoader.asMatrix(imageFile);
        DataNormalization scaler = new ImagePreProcessingScaler(0, 1);
        scaler.transform(image);

        // 进行预测
        INDArray output = model.output(image);
        return Nd4j.argMax(output, 1).getInt(0);
    }
}

import org.deeplearning4j.datasets.iterator.impl.ListDataSetIterator;
import org.deeplearning4j.nn.multilayer.MultiLayerNetwork;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;

import static org.junit.jupiter.api.Assertions.assertNotNull;

public class TrafficSignRecognitionTest {

    private MultiLayerNetwork model;

    @BeforeEach
    public void setup() {
        model = TrafficSignRecognitionModel.buildModel();
    }

    @Test
    public void testLoadData() {
        String dataDirectory = "path/to/your/dataset";
        ListDataSetIterator iterator = DataLoader.loadData(dataDirectory);
        assertNotNull(iterator);
    }

    @Test
    public void testTrainModel() {
        String dataDirectory = "path/to/your/dataset";
        ListDataSetIterator iterator = DataLoader.loadData(dataDirectory);
        TrafficSignRecognitionModel.trainModel(model, iterator);
        assertNotNull(model);
    }
}