diff --git a/modules/python/src2/cv2_convert.cpp b/modules/python/src2/cv2_convert.cpp index b97638064f..e84287cc5a 100644 --- a/modules/python/src2/cv2_convert.cpp +++ b/modules/python/src2/cv2_convert.cpp @@ -151,6 +151,8 @@ bool pyopencv_to(PyObject* o, Mat& m, const ArgInfo& info) bool needcopy = false, needcast = false; int typenum = PyArray_TYPE(oarr), new_typenum = typenum; int type = numpyTypeToCvDepth(typenum); + if( type < 0 ) + type = mlDtypeToCvDepth(typenum); // bfloat16 / float8, when ml_dtypes is installed if( type < 0 ) { @@ -332,12 +334,41 @@ bool pyopencv_to(PyObject* o, Mat& m, const ArgInfo& info) return true; } +// Hands a bfloat16 / float8 Mat back with its exact dtype; `descr` is the caller's to check. +static PyObject* pyopencv_from_mldtype(const cv::Mat& m, PyArray_Descr* descr) +{ + cv::Mat cont; + if( m.isContinuous() ) + cont = m; + else + ERRWRAP2(cont = m.clone()); + + const int cn = cont.channels(); + int dims = cont.dims; + cv::AutoBuffer _sizes(dims + 1); + for( int i = 0; i < dims; i++ ) + _sizes[i] = (npy_intp)cont.size[i]; + if( cn > 1 ) + _sizes[dims++] = cn; + + Py_INCREF(descr); // PyArray_NewFromDescr steals the reference + PyObject* o = PyArray_NewFromDescr(&PyArray_Type, descr, dims, _sizes.data(), + nullptr, nullptr, 0, nullptr); + if( !o ) + return nullptr; + if( !cont.empty() ) + memcpy(PyArray_DATA((PyArrayObject*)o), cont.data, cont.total() * cont.elemSize()); + return o; +} + template<> PyObject* pyopencv_from(const cv::Mat& m) { - // NumPy has no bfloat16 or float8 dtype: widen these to float32 (lossless). if( m.depth() == CV_16BF || m.depth() == CV_8F_E4M3FN || m.depth() == CV_8F_E4M3FNUZ ) { + if( PyArray_Descr* descr = mlDtypeDescrForCvDepth(m.depth()) ) + return pyopencv_from_mldtype(m, descr); + // Without ml_dtypes NumPy cannot name these; widen to float32 (lossless). cv::Mat m32f; ERRWRAP2(m.convertTo(m32f, CV_32F)); return pyopencv_from(m32f); diff --git a/modules/python/src2/cv2_numpy.cpp b/modules/python/src2/cv2_numpy.cpp index d30e777285..fafad8c136 100644 --- a/modules/python/src2/cv2_numpy.cpp +++ b/modules/python/src2/cv2_numpy.cpp @@ -34,19 +34,31 @@ UMatData* NumpyAllocator::allocate(int dims0, const int* sizes, int type, void* int depth = CV_MAT_DEPTH(type); int cn = CV_MAT_CN(type); - // Backing a CV_16BF Mat with the float32 buffer cvDepthToNumpyType() asks for - // would pair a 2-byte element step with 4-byte NumPy strides and corrupt the data. - if( depth == CV_16BF ) - CV_Error(Error::StsNotImplemented, - "CV_16BF (bfloat16) arrays cannot be allocated through the NumPy allocator: " - "NumPy has no bfloat16 dtype"); - int typenum = cvDepthToNumpyType(depth); int i, dims = dims0; cv::AutoBuffer _sizes(dims + 1); for( i = 0; i < dims; i++ ) _sizes[i] = sizes[i]; if( cn > 1 ) _sizes[dims++] = cn; + + // bfloat16 / float8 need their ml_dtypes descriptor: cvDepthToNumpyType()'s float32 + // fallback would pair a narrow element step with wide NumPy strides and corrupt the buffer. + PyArray_Descr* exotic = mlDtypeDescrForCvDepth(depth); + if( exotic ) + { + Py_INCREF(exotic); // PyArray_NewFromDescr steals the reference + PyObject* eo = PyArray_NewFromDescr(&PyArray_Type, exotic, dims, _sizes.data(), + nullptr, nullptr, 0, nullptr); + if(!eo) + CV_Error_(Error::StsError, ("The numpy array of depth=%d, ndims=%d can not be created", depth, dims)); + return allocate(eo, dims0, sizes, type, step); + } + if( depth == CV_16BF || depth == CV_8F_E4M3FN || depth == CV_8F_E4M3FNUZ ) + CV_Error(Error::StsNotImplemented, + "bfloat16 / float8 arrays cannot be allocated through the NumPy allocator: " + "NumPy has no such dtype and the ml_dtypes package is not installed"); + + int typenum = cvDepthToNumpyType(depth); PyObject* o = PyArray_SimpleNew(dims, _sizes.data(), typenum); if(!o) CV_Error_(Error::StsError, ("The numpy array of typenum=%d, ndims=%d can not be created", typenum, dims)); diff --git a/modules/python/src2/cv2_util.cpp b/modules/python/src2/cv2_util.cpp index e5751045df..9bdfffd377 100644 --- a/modules/python/src2/cv2_util.cpp +++ b/modules/python/src2/cv2_util.cpp @@ -1,3 +1,6 @@ +#define NO_IMPORT_ARRAY // the numpy C-API table is imported once, in cv2.cpp +#define PY_ARRAY_UNIQUE_SYMBOL opencv_ARRAY_API + #include "cv2_util.hpp" #include "opencv2/core.hpp" #include "opencv2/core/utils/configuration.private.hpp" @@ -21,7 +24,7 @@ int cvDepthToNumpyType(int depth) case CV_32F: return NPY_FLOAT; case CV_64F: return NPY_DOUBLE; case CV_16F: return NPY_HALF; - // NumPy has no bfloat16 dtype; pyopencv_from() converts the values to float32. + // NumPy alone has no bfloat16; pyopencv_from() uses ml_dtypes when importable, else float32. case CV_16BF: return NPY_FLOAT; case CV_Bool: return NPY_BOOL; default: @@ -30,6 +33,92 @@ int cvDepthToNumpyType(int depth) } } +namespace { + +// ml_dtypes registers bfloat16 and the OCP float8 formats as NumPy dtypes. When importable, +// Mats of OpenCV's matching types keep their exact dtype instead of widening to float32. +struct MlDtypeRegistry +{ + PyArray_Descr* bf16 = nullptr; + PyArray_Descr* e4m3fn = nullptr; + PyArray_Descr* e4m3fnuz = nullptr; + + MlDtypeRegistry() + { + PyObject* mod = PyImport_ImportModule("ml_dtypes"); + if (!mod) + { + PyErr_Clear(); + return; + } + bf16 = fetch(mod, "bfloat16"); + e4m3fn = fetch(mod, "float8_e4m3fn"); + e4m3fnuz = fetch(mod, "float8_e4m3fnuz"); + Py_DECREF(mod); + } + + // The descriptors are kept for the life of the process, so they are never released. + static PyArray_Descr* fetch(PyObject* mod, const char* name) + { + PyObject* scalarType = PyObject_GetAttrString(mod, name); + if (!scalarType) + { + PyErr_Clear(); + return nullptr; + } + PyArray_Descr* descr = nullptr; + if (!PyArray_DescrConverter(scalarType, &descr)) + { + PyErr_Clear(); + descr = nullptr; + } + Py_DECREF(scalarType); + return descr; + } +}; + +// Guarded by the GIL rather than a function-local static: the import inside the constructor +// can release the GIL, and a thread blocked on a static-init lock while holding it deadlocks. +const MlDtypeRegistry* g_mlDtypes = nullptr; + +const MlDtypeRegistry& mlDtypes() +{ + if (!g_mlDtypes) + { + const MlDtypeRegistry* r = new MlDtypeRegistry(); + if (!g_mlDtypes) // another thread may have finished first while the import ran + g_mlDtypes = r; + } + return *g_mlDtypes; +} + +} // namespace + +PyArray_Descr* mlDtypeDescrForCvDepth(int depth) +{ + if (depth != CV_16BF && depth != CV_8F_E4M3FN && depth != CV_8F_E4M3FNUZ) + return nullptr; + const MlDtypeRegistry& r = mlDtypes(); + switch (depth) + { + case CV_16BF: return r.bf16; + case CV_8F_E4M3FN: return r.e4m3fn; + case CV_8F_E4M3FNUZ: return r.e4m3fnuz; + default: return nullptr; + } +} + +int mlDtypeToCvDepth(int typenum) +{ + if (typenum < NPY_USERDEF) // ml_dtypes types are user-defined NumPy dtypes + return -1; + const MlDtypeRegistry& r = mlDtypes(); + if (r.bf16 && typenum == r.bf16->type_num) return CV_16BF; + if (r.e4m3fn && typenum == r.e4m3fn->type_num) return CV_8F_E4M3FN; + if (r.e4m3fnuz && typenum == r.e4m3fnuz->type_num) return CV_8F_E4M3FNUZ; + return -1; +} + int numpyTypeToCvDepth(int typenum) { // Only canonical NPY_* values may be used as case labels: the fixed-width diff --git a/modules/python/src2/cv2_util.hpp b/modules/python/src2/cv2_util.hpp index 4c68094afa..7a98694e01 100644 --- a/modules/python/src2/cv2_util.hpp +++ b/modules/python/src2/cv2_util.hpp @@ -16,6 +16,9 @@ PyObject* failmsgp(const char *fmt, ...); int cvDepthToNumpyType(int depth); int numpyTypeToCvDepth(int typenum); +// nullptr / -1 when ml_dtypes is not importable or the type has no OpenCV counterpart. +PyArray_Descr* mlDtypeDescrForCvDepth(int depth); +int mlDtypeToCvDepth(int typenum); //======================================================================================================================