1 //
2 // SPDX-License-Identifier: BSD-3-Clause
3 // Copyright Contributors to the OpenEXR Project.
4 //
5 
6 // clang-format off
7 
8 #include <Python.h>
9 #include <boost/python.hpp>
10 #include <PyImath.h>
11 #include <PyImathVec.h>
12 #include <PyImathColor.h>
13 #include <iostream>
14 #include <boost/format.hpp>
15 #define NPY_NO_DEPRECATED_API NPY_1_7_API_VERSION
16 #include <numpy/arrayobject.h>
17 
18 using namespace boost::python;
19 using namespace PyImath;
20 
21 template <typename T>
22 struct Holder
23 {
HolderHolder24     Holder( T &a ) : m_val( a ) {}
CleanupHolder25     static void Cleanup (PyObject *capsule)
26     {
27         Holder* h = static_cast<Holder*> (PyCapsule_GetPointer (capsule, NULL));
28         delete h;
29     }
30     T m_val;
31 };
32 
33 template <typename T>
34 static void
setBaseObject(PyObject * nparr,T & arr)35 setBaseObject (PyObject* nparr, T& arr)
36 {
37     using holder         = Holder<T>;
38 
39     holder* ph = new holder (arr);
40     PyObject* capsule = PyCapsule_New (ph, NULL, holder::Cleanup);
41     PyArray_SetBaseObject ((PyArrayObject*) nparr, capsule);
42 }
43 
44 template <typename T> struct NumpyTypeFromType       { BOOST_STATIC_CONSTANT(int, typeEnum=NPY_NOTYPE); };
45 template <> struct NumpyTypeFromType<signed char>    { BOOST_STATIC_CONSTANT(int, typeEnum=NPY_INT8);   };
46 template <> struct NumpyTypeFromType<unsigned char>  { BOOST_STATIC_CONSTANT(int, typeEnum=NPY_UINT8);  };
47 template <> struct NumpyTypeFromType<short>          { BOOST_STATIC_CONSTANT(int, typeEnum=NPY_INT16);  };
48 template <> struct NumpyTypeFromType<unsigned short> { BOOST_STATIC_CONSTANT(int, typeEnum=NPY_UINT16); };
49 template <> struct NumpyTypeFromType<int>            { BOOST_STATIC_CONSTANT(int, typeEnum=NPY_INT32);  };
50 template <> struct NumpyTypeFromType<unsigned int>   { BOOST_STATIC_CONSTANT(int, typeEnum=NPY_UINT32); };
51 template <> struct NumpyTypeFromType<float>          { BOOST_STATIC_CONSTANT(int, typeEnum=NPY_FLOAT);  };
52 template <> struct NumpyTypeFromType<double>         { BOOST_STATIC_CONSTANT(int, typeEnum=NPY_DOUBLE); };
53 
54 template <typename T> struct BaseTypeFrom2DArray       { typedef void type; };
55 template <> struct BaseTypeFrom2DArray<IntArray2D>     { typedef int type; };
56 template <> struct BaseTypeFrom2DArray<FloatArray2D>   { typedef float type; };
57 template <> struct BaseTypeFrom2DArray<DoubleArray2D>  { typedef double type; };
58 template <> struct BaseTypeFrom2DArray<Color4cArray>   { typedef IMATH_NAMESPACE::Color4c type; };
59 template <> struct BaseTypeFrom2DArray<Color4fArray>   { typedef IMATH_NAMESPACE::Color4f type; };
60 
61 template <typename T>
62 object
arrayToNumpy_scalar(T & sa)63 arrayToNumpy_scalar(T &sa)
64 {
65     typedef typename T::BaseType PodType;
66 
67     if (sa.stride() != 1)
68         throw std::logic_error("Unable to make numpy wrapping of strided arrays");
69 
70     int type = NumpyTypeFromType<PodType>::typeEnum;
71     npy_intp dims = sa.len();
72     PodType *data = &sa[0];
73     PyObject *a = PyArray_SimpleNewFromData(1, &dims, type, data);
74 
75     if (!a)
76         throw_error_already_set();
77 
78     setBaseObject (a, sa);
79 
80     object retval = object(handle<>(a));
81     return retval;
82 }
83 
84 template<typename T>
85 object
arrayToNumpy_vector(T & va)86 arrayToNumpy_vector(T &va)
87 {
88     typedef typename T::BaseType        BaseType;
89     typedef typename BaseType::BaseType PodType;
90 
91     if (va.stride() != 1)
92         throw std::logic_error("Unable to make numpy wrapping of strided arrays");
93 
94     int type = NumpyTypeFromType<PodType>::typeEnum;
95     npy_intp dims[2]{ va.len(), BaseType::dimensions()};
96     PodType *data = &va[0][0];
97     PyObject *a = PyArray_SimpleNewFromData(2, dims, type, data);
98 
99     if (!a)
100         throw_error_already_set();
101 
102     setBaseObject (a, va);
103 
104     object retval = object (handle<> (a));
105     return retval;
106 }
107 
108 template<typename T>
109 object
arrayToNumpy_scalar2D(T & ca)110 arrayToNumpy_scalar2D(T &ca)
111 {
112     typedef typename BaseTypeFrom2DArray<T>::type PodType;
113 
114     // Numpy is matrix indexed based - the first index indicates the row and
115     // the second index indicates the column.
116     // Zeno data is image indexed based - the first index indicates the column
117     // and the second indicates the row.
118     // Swap the dimensions so Numpy handles data natively with matrix based
119     // indexing.
120     IMATH_NAMESPACE::Vec2<size_t> len = ca.len();
121     int type = NumpyTypeFromType<PodType>::typeEnum;
122     npy_intp dims[2]{ static_cast<npy_intp>(len.y),
123                       static_cast<npy_intp>(len.x) };
124     PodType *data = &ca(0, 0);
125     PyObject *a = PyArray_SimpleNewFromData(2, dims, type, data);
126 
127     if (!a)
128         throw_error_already_set();
129 
130     setBaseObject (a, ca);
131 
132     object retval = object(handle<>(a));
133     return retval;
134 }
135 
136 template<typename T>
137 object
arrayToNumpy_vector2D(T & ca)138 arrayToNumpy_vector2D(T &ca)
139 {
140     typedef typename BaseTypeFrom2DArray<T>::type BaseType;
141     typedef typename BaseType::BaseType           PodType;
142 
143     // Numpy is matrix indexed based - the first index indicates the row and
144     // the second index indicates the column.
145     // Zeno data is image indexed based - the first index indicates the column
146     // and the second indicates the row.
147     // Swap the dimensions so Numpy handles data natively with matrix based
148     // indexing.
149     IMATH_NAMESPACE::Vec2<size_t> len = ca.len();
150     int type = NumpyTypeFromType<PodType>::typeEnum;
151     npy_intp dims[3]{ static_cast<npy_intp>(len.y),
152                       static_cast<npy_intp>(len.x),
153                       BaseType::dimensions() };
154     PodType *data = &ca(0, 0)[0];
155     PyObject *a = PyArray_SimpleNewFromData(3, dims, type, data);
156 
157     if (!a)
158         throw_error_already_set();
159 
160     setBaseObject (a, ca);
161 
162     object retval = object(handle<>(a));
163     return retval;
164 }
165 
166 #if PY_MAJOR_VERSION > 2
apply_import()167 static void *apply_import()
168 {
169     import_array();
170     return 0;
171 }
172 #endif
173 
174 // Convenience macro for wrapping scalar PyImath array types.
175 #define WRAP_SCALAR_ARRAY(ARRAY_TYPE)                                            \
176     def("arrayToNumpy", &arrayToNumpy_scalar<ARRAY_TYPE>,                        \
177         "arrayToNumpy(array) - wrap the given " #ARRAY_TYPE " as a numpy array", \
178         (arg("array")));
179 #define WRAP_SCALAR_ARRAY_2D(ARRAY_TYPE)                                         \
180     def("arrayToNumpy", &arrayToNumpy_scalar2D<ARRAY_TYPE>,                      \
181         "arrayToNumpy(array) - wrap the given " #ARRAY_TYPE " as a numpy array", \
182         (arg("array")));
183 
184 // Convenience macro for wrapping vector PyImath array types.
185 #define WRAP_VECTOR_ARRAY(ARRAY_TYPE)                                            \
186     def("arrayToNumpy", &arrayToNumpy_vector<ARRAY_TYPE>,                        \
187         "arrayToNumpy(array) - wrap the given " #ARRAY_TYPE " as a numpy array", \
188         (arg("array")));
189 #define WRAP_VECTOR_ARRAY_2D(ARRAY_TYPE)                                         \
190     def("arrayToNumpy", &arrayToNumpy_vector2D<ARRAY_TYPE>,                      \
191         "arrayToNumpy(array) - wrap the given " #ARRAY_TYPE " as a numpy array", \
192         (arg("array")));
193 
BOOST_PYTHON_MODULE(imathnumpy)194 BOOST_PYTHON_MODULE(imathnumpy)
195 {
196     scope().attr("__doc__") = "Imathnumpy module";
197     scope().attr("__version__") = IMATH_VERSION_STRING;
198 
199     handle<> imath(PyImport_ImportModule("imath"));
200     if (PyErr_Occurred()) throw_error_already_set();
201     scope().attr("imath") = imath;
202 
203     handle<> numpy(PyImport_ImportModule("numpy"));
204     if (PyErr_Occurred()) throw_error_already_set();
205     scope().attr("numpy") = numpy;
206 
207 #if PY_MAJOR_VERSION > 2
208     // seems like numpy expects this to be used in a scenario
209     // where there is a return value in python3...
210     (void)apply_import();
211 #else
212     import_array();
213 #endif
214 
215     scope().attr("__doc__") = "Array wrapping module to overlay imath array data with numpy arrays";
216 
217     WRAP_SCALAR_ARRAY(SignedCharArray)
218     WRAP_SCALAR_ARRAY(UnsignedCharArray)
219     WRAP_SCALAR_ARRAY(ShortArray)
220     WRAP_SCALAR_ARRAY(UnsignedShortArray)
221     WRAP_SCALAR_ARRAY(IntArray)
222     WRAP_SCALAR_ARRAY(UnsignedIntArray)
223     WRAP_SCALAR_ARRAY(FloatArray)
224     WRAP_SCALAR_ARRAY(DoubleArray)
225 
226     WRAP_VECTOR_ARRAY(V2sArray)
227     WRAP_VECTOR_ARRAY(V2iArray)
228     WRAP_VECTOR_ARRAY(V2fArray)
229     WRAP_VECTOR_ARRAY(V2dArray)
230 
231     WRAP_VECTOR_ARRAY(V3sArray)
232     WRAP_VECTOR_ARRAY(V3iArray)
233     WRAP_VECTOR_ARRAY(V3fArray)
234     WRAP_VECTOR_ARRAY(V3dArray)
235 
236     WRAP_VECTOR_ARRAY(V4sArray)
237     WRAP_VECTOR_ARRAY(V4iArray)
238     WRAP_VECTOR_ARRAY(V4fArray)
239     WRAP_VECTOR_ARRAY(V4dArray)
240 
241     WRAP_VECTOR_ARRAY(C3cArray)
242     WRAP_VECTOR_ARRAY(C3fArray)
243     WRAP_VECTOR_ARRAY(C4cArray)
244     WRAP_VECTOR_ARRAY(C4fArray)
245 
246     WRAP_SCALAR_ARRAY_2D(IntArray2D)
247     WRAP_SCALAR_ARRAY_2D(FloatArray2D)
248     WRAP_SCALAR_ARRAY_2D(DoubleArray2D)
249 
250     WRAP_VECTOR_ARRAY_2D(Color4cArray)
251     WRAP_VECTOR_ARRAY_2D(Color4fArray)
252 }
253