Java中如何直接将BufferedImage转换为Uint8 Tensor用于TF2推理

无需通过磁盘读写JPEG文件，直接从BufferedImage提取像素数据构造Tensor即可，核心步骤如下：

1. 提取BufferedImage的像素信息，转换为模型需要的通道格式（通常为RGB）
2. 将像素数据整理为NHWC维度的数组（[batch_size, height, width, channels]）
3. 利用TensorFlow Java API直接创建TUint8类型的Tensor

完整代码示例 #

SavedModelBundle model = SavedModelBundle.load(modelPath,"serve");
BufferedImage drawImg = BIfromMat(aMat);

// 获取图像尺寸参数
int height = drawImg.getHeight();
int width = drawImg.getWidth();
int channels = 3; // 对应RGB通道
long[] tensorShape = {1, height, width, channels}; // batch=1，符合模型输入的4D格式

// 提取并转换像素数据：将BufferedImage的ARGB格式转为RGB字节数组
byte[] pixelBytes = new byte[height * width * channels];
int index = 0;
for (int y = 0; y < height; y++) {
    for (int x = 0; x < width; x++) {
        int argbPixel = drawImg.getRGB(x, y);
        // 按RGB顺序提取通道值（若模型要求BGR则调换顺序为B、G、R）
        pixelBytes[index++] = (byte) ((argbPixel >> 16) & 0xFF); // R通道
        pixelBytes[index++] = (byte) ((argbPixel >> 8) & 0xFF);  // G通道
        pixelBytes[index++] = (byte) (argbPixel & 0xFF);         // B通道
    }
}

// 创建Tensor并执行推理
try (TUint8 inputTensor = TUint8.tensorOf(Shape.make(tensorShape), ByteBuffer.wrap(pixelBytes))) {
    Map<String, Tensor> feedDict = new HashMap<>();
    feedDict.put("input_tensor", inputTensor);
    Map<String, Tensor> outputTensorMap = model.function("serving_default").call(feedDict);
    
    // 在这里处理检测结果...
}

注意事项 #

• 若你的模型输入要求BGR格式（比如直接沿用OpenCV的通道顺序），只需调整像素提取的顺序为B、G、R
• 如果是灰度图像，将channels设为1，并修改像素提取逻辑只取单通道值
• 该方式完全规避磁盘IO，推理效率会大幅提升

内容的提问来源于stack exchange，提问作者Robert Huber