利用深度学习实现验证码识别-2-使用Python导出ONNX模型并在Java中调用实现验证码识别

时间：09-24来源：作者：点击数：47

1. Python部分：导出ONNX模型

首先，我们需要在Python中定义并导出一个已经训练好的验证码识别模型。以下是完整的Python代码：

import string
import torch
import torch.nn as nn
import torch.nn.functional as F

CHAR_SET = string.digits

# 优化后的模型设计
class CaptchaModel(nn.Module):
    def __init__(self):
        super(CaptchaModel, self).__init__()
        self.conv1 = nn.Conv2d(1, 32, 3, padding=1)
        self.conv2 = nn.Conv2d(32, 64, 3, padding=1)
        self.conv3 = nn.Conv2d(64, 128, 3, padding=1)
        self.fc1 = nn.Linear(128 * 5 * 12, 256)  # 调整为实际展平维度
        self.fc2 = nn.Linear(256, 4 * len(CHAR_SET))
        self.dropout = nn.Dropout(0.5)

    def forward(self, x):
        x = F.relu(F.max_pool2d(self.conv1(x), 2))
        x = F.relu(F.max_pool2d(self.conv2(x), 2))
        x = F.relu(F.max_pool2d(self.conv3(x), 2))
        x = x.view(x.size(0), -1)
        x = F.relu(self.fc1(x))
        x = self.dropout(x)
        x = self.fc2(x)
        return x.view(-1, 4, len(CHAR_SET))

# 使用CUDA，如果可用的话
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Using device: {device}")
# 假设你的模型已经训练好并保存在 'best_model.pth'
model = CaptchaModel().to(device)
model.load_state_dict(torch.load('best_model.pth'))

# 生成一个测试输入 (示例输入的形状应与模型输入形状一致)
dummy_input = torch.randn(1, 1, 40, 100).to(device)

# 导出模型为 ONNX 格式
torch.onnx.export(model, dummy_input, "captcha_model.onnx", 
                  input_names=["input"], output_names=["output"], 
                  dynamic_axes={"input": {0: "batch_size"}, "output": {0: "batch_size"}})

print("Model exported to captcha_model.onnx")

这段代码定义了一个验证码识别模型，并将其导出为ONNX格式，以便在Java中使用。

2. Java部分：调用ONNX模型进行验证码识别

接下来，我们使用Java调用导出的ONNX模型进行验证码识别。以下是完整的Java代码：

引用onnxruntime-1.19.0.jar

package com.tushuoit;

import ai.onnxruntime.*;
import javax.imageio.ImageIO;
import java.awt.*;
import java.awt.image.BufferedImage;
import java.io.File;
import java.nio.FloatBuffer;
import java.util.ArrayList;
import java.util.Collections;
import java.util.Random;
import java.util.List;

public class CaptchaInference {
    private static final String CHAR_SET = "0123456789";
    private static final int INPUT_WIDTH = 100;
    private static final int INPUT_HEIGHT = 40;
    private static final Random random = new Random();

    public static void main(String[] args) throws Exception {
        // 随机生成4个字符的验证码文本
        String captchaText = generateRandomText(4);
        System.out.println("Generated Captcha Text: " + captchaText);

        // 生成包含文本的Bitmap (BufferedImage)
        BufferedImage captchaImage = generateCaptcha(captchaText, 36, INPUT_WIDTH, INPUT_HEIGHT);

        // 将Bitmap保存为文件（仅用于查看生成的图像，实际使用中可以省略）
        ImageIO.write(captchaImage, "png", new File("generated_captcha.png"));

        // 将图像转换为浮点数数组，并进行归一化处理
        float[] inputData = imageToFloatArray(captchaImage);

        // 创建ONNX Runtime环境
        OrtEnvironment env = OrtEnvironment.getEnvironment();
        OrtSession.SessionOptions opts = new OrtSession.SessionOptions();

        // 加载ONNX模型
        OrtSession session = env.createSession("captcha_model.onnx", opts);

        // 创建输入张量
        FloatBuffer inputBuffer = FloatBuffer.wrap(inputData);
        OnnxTensor inputTensor = OnnxTensor.createTensor(env, inputBuffer,
                new long[] { 1, 1, INPUT_HEIGHT, INPUT_WIDTH });

        // 进行推理
        OrtSession.Result result = session.run(Collections.singletonMap("input", inputTensor));

        // Extract output tensor and decode it
        float[][][] outputData = (float[][][]) result.get(0).getValue();
        List<String> decodedTexts = decodeOutput(outputData);

        // Print the decoded captcha text
        for (String text : decodedTexts) {
            System.out.println("Predicted Captcha Text: " + text);
        }

        System.out.println("Inference completed.");
        // 释放资源
        session.close();
        env.close();
    }

    // 随机生成指定长度的验证码文本
    private static String generateRandomText(int length) {
        StringBuilder text = new StringBuilder(length);
        for (int i = 0; i < length; i++) {
            text.append(CHAR_SET.charAt(random.nextInt(CHAR_SET.length())));
        }
        return text.toString();
    }

    // 生成包含文本的BufferedImage
    private static BufferedImage generateCaptcha(String text, int fontSize, int width, int height) {
        BufferedImage image = new BufferedImage(width, height, BufferedImage.TYPE_INT_RGB);
        Graphics2D g2d = image.createGraphics();

        // 设置背景颜色为白色
        g2d.setColor(Color.WHITE);
        g2d.fillRect(0, 0, width, height);

        // 设置字体和颜色
        g2d.setFont(new Font("DroidSansMono", Font.PLAIN, fontSize));
        g2d.setColor(Color.BLACK);

        // 绘制文本
        FontMetrics fm = g2d.getFontMetrics();
        int x = 5; // 文字开始的X坐标
        int y = fm.getAscent() + 5; // 文字开始的Y坐标
        g2d.drawString(text, x, y);

        g2d.dispose();
        return image;
    }

    // 将BufferedImage转换为float数组，并进行归一化处理
    private static float[] imageToFloatArray(BufferedImage image) {
        int width = image.getWidth();
        int height = image.getHeight();
        float[] floatArray = new float[width * height];

        for (int y = 0; y < height; y++) {
            for (int x = 0; x < width; x++) {
                int rgb = image.getRGB(x, y);
                int gray = (rgb >> 16) & 0xFF; // 因为是灰度图，只需获取一个通道的值
                floatArray[y * width + x] = (gray / 255.0f - 0.5f) * 2.0f; // 归一化到[-1, 1]
            }
        }

        return floatArray;
    }

    private static List<String> decodeOutput(float[][][] outputData) {
        List<String> decodedTexts = new ArrayList<>();
        for (float[][] singleOutput : outputData) {
            StringBuilder decodedText = new StringBuilder();
            for (float[] charProbabilities : singleOutput) {
                int maxIndex = getMaxIndex(charProbabilities);
                decodedText.append(CHAR_SET.charAt(maxIndex));
            }
            decodedTexts.add(decodedText.toString());
        }
        return decodedTexts;
    }

    private static int getMaxIndex(float[] probabilities) {
        int maxIndex = 0;
        float maxProb = probabilities[0];
        for (int i = 1; i < probabilities.length; i++) {
            if (probabilities[i] > maxProb) {
                maxProb = probabilities[i];
                maxIndex = i;
            }
        }
        return maxIndex;
    }
}