numpy_to_mat.cpp
Go to the documentation of this file.
00001 #include <boost/python.hpp>
00002 #include <numpy/arrayobject.h>
00003 #include <complex>
00004 #include <algorithm>
00005 
00006 #include <opencv2/core/core.hpp>
00007 #include <pano_py/opencv.h>
00008 
00009 using namespace cv;
00010 
00011 namespace{
00012 
00013 static size_t REFCOUNT_OFFSET = (size_t)&(((PyObject*)0)->ob_refcnt) +
00014     (0x12345678 != *(const size_t*)"\x78\x56\x34\x12\0\0\0\0\0")*sizeof(int);
00015 
00016 static inline PyObject* pyObjectFromRefcount(const int* refcount)
00017 {
00018     return (PyObject*)((size_t)refcount - REFCOUNT_OFFSET);
00019 }
00020 
00021 static inline int* refcountFromPyObject(const PyObject* obj)
00022 {
00023     return (int*)((size_t)obj + REFCOUNT_OFFSET);
00024 }
00025 
00026 class NumpyAllocator : public MatAllocator
00027 {
00028 public:
00029     NumpyAllocator(){}
00030     ~NumpyAllocator() {}
00031 
00032     void allocate(int dims, const int* sizes, int type, int*& refcount,
00033                   uchar*& datastart, uchar*& data, size_t* step)
00034     {
00035         int depth = CV_MAT_DEPTH(type);
00036         int cn = CV_MAT_CN(type);
00037         const int f = (int)(sizeof(size_t)/8);
00038         int typenum = depth == CV_8U ? NPY_UBYTE : depth == CV_8S ? NPY_BYTE :
00039                       depth == CV_16U ? NPY_USHORT : depth == CV_16S ? NPY_SHORT :
00040                       depth == CV_32S ? NPY_INT : depth == CV_32F ? NPY_FLOAT :
00041                       depth == CV_64F ? NPY_DOUBLE : f*NPY_ULONGLONG + (f^1)*NPY_UINT;
00042 
00043         npy_intp _sizes[CV_MAX_DIM+1];
00044         for( int i = 0; i < dims; i++ )
00045             _sizes[i] = sizes[i];
00046         if( cn > 1 )        {
00047             if( _sizes[dims-1] == 1 )
00048                 _sizes[dims-1] = cn;
00049             else
00050                 _sizes[dims++] = cn;
00051         }
00052 
00053         PyObject* o = PyArray_SimpleNew(dims, _sizes, typenum);
00054         if(!o)
00055             CV_Error_(CV_StsError, ("The numpy array of typenum=%d, ndims=%d can not be created", typenum, dims));
00056         refcount = refcountFromPyObject(o);
00057         npy_intp* _strides = PyArray_STRIDES(o);
00058         for(int i = 0; i < dims - (cn > 1); i++ )
00059             step[i] = (size_t)_strides[i];
00060         datastart = data = (uchar*)PyArray_DATA(o);
00061     }
00062 
00063     void deallocate(int* refcount, uchar* datastart, uchar* data)
00064     {
00065         if( !refcount )
00066             return;
00067         PyObject* o = pyObjectFromRefcount(refcount);
00068         Py_INCREF(o);
00069         Py_DECREF(o);
00070     }
00071 };
00072 
00073 NumpyAllocator g_numpyAllocator;
00074 
00075 } //end anon namespace
00076 
00077 namespace pano_py{
00078 
00079 bool numpy_to_mat(const PyObject* o, cv::Mat& m, const char* name, bool allowND)
00080 {
00081   if(!o || o == Py_None)
00082     {
00083         if( !m.data )
00084             m.allocator = &g_numpyAllocator;
00085         return true;
00086     }
00087 
00088     if( !PyArray_Check(o) )
00089     {
00090         failmsg("%s is not a numpy array", name);
00091         return false;
00092     }
00093 
00094     int typenum = PyArray_TYPE(o);
00095     int type = typenum == NPY_UBYTE ? CV_8U : typenum == NPY_BYTE ? CV_8S :
00096                typenum == NPY_USHORT ? CV_16U : typenum == NPY_SHORT ? CV_16S :
00097                typenum == NPY_INT || typenum == NPY_LONG ? CV_32S :
00098                typenum == NPY_FLOAT ? CV_32F :
00099                typenum == NPY_DOUBLE ? CV_64F : -1;
00100 
00101     if( type < 0 )
00102     {
00103         failmsg("%s data type = %d is not supported", name, typenum);
00104         return false;
00105     }
00106 
00107     int ndims = PyArray_NDIM(o);
00108     if(ndims >= CV_MAX_DIM)
00109     {
00110         failmsg("%s dimensionality (=%d) is too high", name, ndims);
00111         return false;
00112     }
00113 
00114     int size[CV_MAX_DIM+1];
00115     size_t step[CV_MAX_DIM+1], elemsize = CV_ELEM_SIZE1(type);
00116     const npy_intp* _sizes = PyArray_DIMS(o);
00117     const npy_intp* _strides = PyArray_STRIDES(o);
00118 
00119     for(int i = 0; i < ndims; i++)
00120     {
00121         size[i] = (int)_sizes[i];
00122         step[i] = (size_t)_strides[i];
00123     }
00124 
00125     if( ndims == 0 || step[ndims-1] > elemsize ) {
00126         size[ndims] = 1;
00127         step[ndims] = elemsize;
00128         ndims++;
00129     }
00130 
00131     if( ndims == 3 && size[2] <= CV_CN_MAX && step[1] == elemsize*size[2] )
00132     {
00133         ndims--;
00134         type |= CV_MAKETYPE(0, size[2]);
00135     }
00136 
00137     if( ndims > 2 && !allowND )
00138     {
00139         failmsg("%s has more than 2 dimensions", name);
00140         return false;
00141     }
00142 
00143     m = Mat(ndims, size, type, PyArray_DATA(o), step);
00144 
00145     if( m.data )
00146     {
00147         m.refcount = refcountFromPyObject(o);
00148         m.addref(); // protect the original numpy array from deallocation
00149                     // (since Mat destructor will decrement the reference counter)
00150     };
00151     m.allocator = &g_numpyAllocator;
00152     return true;
00153 }
00154 } //namespace pano_py
00155 


pano_py
Author(s): Ethan Rublee
autogenerated on Wed Aug 26 2015 16:34:19