使用distutils封装C++ OpenCV扩展为Python库时，如何返回cv::Mat数据？

我完全懂你现在的困扰——用distutils封装C++ OpenCV代码时，不知道怎么把cv::Mat转换成Python能识别的对象对吧？其实核心思路是把cv::Mat转换成numpy数组，因为Python版OpenCV本身就是用numpy数组来表示图像的，这样返回的结果能直接和Python的cv2库无缝配合。

下面是完整的修改方案，一步步来：

1. 补充必要的头文件和初始化代码 #

首先，我们需要引入numpy的C API来处理数组转换，还要在模块初始化时初始化numpy的接口：

#include <Python.h>
#include <opencv2/core.hpp>
#include <opencv2/imgcodecs.hpp>
#include <numpy/arrayobject.h>  // 新增：numpy的C API头文件

using namespace cv;

// 新增：把cv::Mat转换成numpy数组的工具函数
static PyObject* mat_to_numpy(const cv::Mat& mat) {
    // 映射OpenCV的图像类型到numpy的dtype
    int numpy_type;
    switch (mat.depth()) {
        case CV_8U: numpy_type = NPY_UINT8; break;
        case CV_16U: numpy_type = NPY_UINT16; break;
        case CV_32S: numpy_type = NPY_INT32; break;
        case CV_32F: numpy_type = NPY_FLOAT32; break;
        case CV_64F: numpy_type = NPY_FLOAT64; break;
        default:
            PyErr_SetString(PyExc_TypeError, "Unsupported image depth type");
            return nullptr;
    }

    // 确定numpy数组的维度：单通道是2维(HxW)，多通道是3维(HxWxC)
    npy_intp dims[3];
    int ndim;
    if (mat.channels() == 1) {
        dims[0] = mat.rows;
        dims[1] = mat.cols;
        ndim = 2;
    } else {
        dims[0] = mat.rows;
        dims[1] = mat.cols;
        dims[2] = mat.channels();
        ndim = 3;
    }

    // 创建空的numpy数组，然后复制Mat的数据进去
    PyObject* numpy_arr = PyArray_SimpleNew(ndim, dims, numpy_type);
    if (!numpy_arr) return nullptr;

    // 复制数据：避免C++的Mat销毁后numpy数组出现悬空指针
    void* numpy_data = PyArray_DATA((PyArrayObject*)numpy_arr);
    memcpy(numpy_data, mat.data, mat.total() * mat.elemSize());

    return numpy_arr;
}

// 修改你的imread函数：支持从Python传入文件名，处理错误，返回numpy数组
static PyObject* cv_imread(PyObject *self, PyObject *args) {
    const char* img_path;
    // 解析Python传入的参数（这里是字符串类型的文件名）
    if (!PyArg_ParseTuple(args, "s", &img_path)) {
        return nullptr;  // 参数解析失败时自动设置Python异常
    }

    cv::Mat img = cv::imread(img_path, cv::IMREAD_COLOR);
    if (img.empty()) {
        PyErr_SetString(PyExc_IOError, "Failed to read the image file");
        return nullptr;
    }

    // 转换并返回numpy数组
    return mat_to_numpy(img);
}

// 模块方法定义
static PyMethodDef CvConverterMethods[] = {
    {"imread", cv_imread, METH_VARARGS, "Read image from path and return as numpy array"},
    {nullptr, nullptr, 0, nullptr}  // 哨兵元素，标记方法列表结束
};

// 模块结构定义
static struct PyModuleDef cv_converter_module = {
    PyModuleDef_HEAD_INIT,
    "cv_converter",  // 模块名，Python中import时用这个名字
    "OpenCV-Python image conversion module",  // 模块描述
    -1,  // 不需要per-interpreter状态
    CvConverterMethods
};

// 模块初始化函数：必须命名为PyInit_<模块名>
PyMODINIT_FUNC PyInit_cv_converter(void) {
    PyObject* m = PyModule_Create(&cv_converter_module);
    if (!m) return nullptr;

    // 初始化numpy的C API，必须在模块创建后调用
    import_array();
    return m;
}

2. 编写distutils的setup.py #

接下来写setup.py，告诉distutils怎么编译这个扩展，注意要包含numpy和OpenCV的头文件、库文件路径：

from distutils.core import setup, Extension
import numpy as np

# 替换成你的OpenCV头文件和库文件路径
OPENCV_INCLUDE_DIR = "/usr/local/include/opencv4"  # 示例路径，根据你的安装位置调整
OPENCV_LIB_DIR = "/usr/local/lib"  # 示例路径

cv_converter_ext = Extension(
    "cv_converter",
    sources=["cv_converter.cpp"],
    include_dirs=[np.get_include(), OPENCV_INCLUDE_DIR],
    libraries=["opencv_core", "opencv_imgcodecs"],
    library_dirs=[OPENCV_LIB_DIR]
)

setup(
    name="cv_converter",
    version="1.0",
    description="A C++ OpenCV extension for Python",
    ext_modules=[cv_converter_ext]
)

3. 编译和测试 #

编译扩展：

python setup.py build_ext --inplace

然后在Python中测试：

import cv_converter
import cv2

# 调用扩展的imread函数
img = cv_converter.imread("sample.png")
print(img.shape)  # 输出图像的形状，比如(480, 640, 3)

# 可以直接用cv2显示图像
cv2.imshow("Image", img)
cv2.waitKey(0)
cv2.destroyAllWindows()

关键注意点 #

• 内存安全：上面的代码用了memcpy复制cv::Mat的数据到numpy数组，这样C++的Mat销毁后，numpy数组依然持有独立的数据，避免悬空指针问题。如果想优化内存使用，可以尝试共享内存，但需要额外处理内存生命周期，新手建议先使用复制方案。
• 通道顺序：OpenCV默认是BGR通道顺序，和Python版cv2的imread返回的数组一致，所以不需要额外转换。如果需要RGB，可以在C++里用cv::cvtColor(img, img, COLOR_BGR2RGB)再转换。
• 错误处理：一定要设置Python异常（比如PyErr_SetString），这样Python端能捕获到错误信息，方便调试。

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