Golang中如何通过gocv将图像转换为TensorFlow Serving可用张量

需要将对接TensorFlow Serving的Python图像预处理逻辑转换为Golang实现，原Python逻辑如下：

image = cv2.imread("1.jpg")
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
image = cv2.resize(image, (256, 256))
image = np.expand_dims(image, axis=0)
image = np.true_divide(image, 255)
np.shape(image)

已完成初步Gocv代码编写，但存在三个待解决问题：

• 无法实现np.expand_dims对应的维度扩展操作
• 无法实现像素值除以255的归一化操作
• 尝试对接gonum库未成功，无法生成符合TensorFlow Serving输入要求的张量
  原有初步代码如下：

package main

import (
    "log"
    "image"
    "gocv.io/x/gocv"
)
func main() {
    imageFilePath := "./a.jpg"
    mat := gocv.IMRead(imageFilePath, gocv.IMReadAnyColor)
    if mat.Empty() {
        log.Panic("Can not read Image file : ", imageFilePath)
        return
    } 
    resizeImage := gocv.NewMat()
    gocv.Resize(mat, &resizeImage, image.Point{X: 256, Y: 256}, 0, 0, gocv.InterpolationNearestNeighbor)
 
    img := resizeImage.Clone()
    gocv.CvtColor(resizeImage, &img, gocv.ColorBGRToRGB)
    // equivalent to np.expand_dims(image, axis=0) ? 

}

不需要引入gonum等额外矩阵库，直接基于Gocv即可完成全部预处理逻辑，核心注意点如下：

• np.expand_dims(axis=0)仅在矩阵最外层增加长度为1的batch维度，不会改变像素在内存中的排列顺序，无需修改实际像素数据，只需要在构造TensorFlow张量时指定shape为[1,256,256,3]即可实现等价效果
• 像素归一化可以直接通过Gocv的ConvertTo方法一步完成，支持在类型转换的同时传入缩放系数，性能远高于手动遍历像素
• 原代码中resize使用的InterpolationNearestNeighbor和OpenCV Python默认的双线性插值行为不一致，需要替换为InterpolationLinear保证预处理结果对齐
• Gocv的Mat需要手动调用Close()释放内存，避免泄漏

完整可运行代码如下：

package main

import (
	"log"
	"unsafe"

	"gocv.io/x/gocv"
	tf "github.com/tensorflow/tensorflow/tensorflow/go"
)

func main() {
	imageFilePath := "./a.jpg"
	// 对应Python: image = cv2.imread("1.jpg")
	mat := gocv.IMRead(imageFilePath, gocv.IMReadColor)
	if mat.Empty() {
		log.Panic("failed to read image: ", imageFilePath)
	}
	defer mat.Close()

	// 对应Python: image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
	rgbMat := gocv.NewMat()
	defer rgbMat.Close()
	gocv.CvtColor(mat, &rgbMat, gocv.ColorBGRToRGB)

	// 对应Python: image = cv2.resize(image, (256, 256))
	// 使用InterpolationLinear对齐Python OpenCV默认插值方式
	resizedMat := gocv.NewMat()
	defer resizedMat.Close()
	gocv.Resize(rgbMat, &resizedMat, image.Pt(256, 256), 0, 0, gocv.InterpolationLinear)

	// 对应Python: image = np.true_divide(image, 255)
	// 一步完成uint8转float32 + 像素值除以255归一化
	normalizedMat := gocv.NewMat()
	defer normalizedMat.Close()
	resizedMat.ConvertTo(&normalizedMat, gocv.MatTypeCV32FC3, 1.0/255.0, 0)

	// 提取Mat中连续存储的float32像素数据
	// 内存排列顺序为(H,W,C)即(256,256,3)，和Python NumPy数组内存布局完全一致
	imgBytes := normalizedMat.ToBytes()
	var tensorData []float32
	// 零拷贝将字节流转换为float32切片
	tensorHeader := (*[2]unsafe.Pointer)(unsafe.Pointer(&tensorData))
	bytesHeader := (*[2]unsafe.Pointer)(unsafe.Pointer(&imgBytes))
	tensorHeader[0] = bytesHeader[0]
	// 切片长度：总字节数 / 每个float32占4字节
	tensorHeader[1] = unsafe.Pointer(uintptr(bytesHeader[1]) / 4)

	// 对应Python: image = np.expand_dims(image, axis=0)
	// 构造TF张量，指定shape为(1,256,256,3)，直接传入提取的float32数据即可
	inputTensor, err := tf.NewTensor(tensorData, []int64{1, 256, 256, 3})
	if err != nil {
		log.Panic("failed to create input tensor: ", err)
	}

	// 此处即可将inputTensor传入TensorFlow Serving客户端完成推理
	log.Printf("tensor created successfully, shape: %v", inputTensor.Shape())
}

验证说明 #

• 最终生成的张量数据长度为1*256*256*3 = 196608，和Python端np.shape(image)输出的(1,256,256,3)完全对齐
• 所有预处理操作均调用OpenCV原生实现，性能和Python版OpenCV一致，无额外性能损耗
• 如果不想使用unsafe包做零拷贝转换，可以遍历字节流逐个解析float32值，性能略低但兼容性更好；如果使用gRPC方式调用TF Serving，直接将float32切片按部署环境的字节序序列化后传入即可。

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