Python: ImageBuf constructor/reset from numpy array

src/doc/pythonbindings.rst

@@ -1556,6 +1556,28 @@ awaiting a call to `reset()` or `copy()` before it is useful.
         buf = ImageBuf (spec)
 
 
+.. py:method:: ImageBuf (data)
+
+    Construct a writable ImageBuf of the dimensions of `data`, which is a
+    NumPy `ndarray` of values indexed as `[y][x][channel]` for normal 2D
+    images, or for 3D volumetric images, as `[z][y][x][channel]`. The data
+    will be copied into the ImageBuf's internal storage. The NumPy array may
+    be strided for z, y, or x, but must have "contiguous" channel data within
+    each pixel. The pixel data type is also deduced from the contents of the
+    `data` array.
+
+    Note that this Python ImageBuf will contain its own copy of the data, so
+    further changes to the `data` array will not affect the ImageBuf. This is
+    different from the C++ ImageBuf constructor from a pointer, which will
+    "wrap" the existing user-provided buffer but not make its own copy.
+
+    Example:
+
+    .. code-block:: python
+
+        pixels = numpy.zeros ((640, 480, 3), dtype = numpy.float32)
+        buf = ImageBuf (pixels)
+
 
 .. py:method:: ImageBuf.clear ()
 
@@ -1589,6 +1611,22 @@ awaiting a call to `reset()` or `copy()` before it is useful.
     uninitialized.
 
 
+.. py:method:: ImageBuf.reset (data)
+
+    Reset the ImageBuf to be sized to the dimensions of `data`, which is a
+    NumPy `ndarray` of values indexed as `[y][x][channel]` for normal 2D
+    images, or for 3D volumetric images, as `[z][y][x][channel]`. The data
+    will be copied into the ImageBuf's internal storage. The NumPy array may
+    be strided for z, y, or x, but must have "contiguous" channel data within
+    each pixel. The pixel data type is also deduced from the contents of the
+    `data` array.
+
+    Note that this Python ImageBuf will contain its own copy of the data, so
+    further changes to the `data` array will not affect the ImageBuf. This is
+    different from the C++ ImageBuf constructor from a pointer, which will
+    "wrap" the existing user-provided buffer but not make its own copy.
+
+
 .. py:method:: ImageBuf.read(subimage=0, miplevel=0, force=False, convert=oiio.UNKNOWN)
                ImageBuf.read(subimage, miplevel, chbegin, chend, force, convert)
 

src/include/OpenImageIO/imagebuf.h

@@ -807,10 +807,12 @@ class OIIO_API ImageBuf {
 
     /// Copy the data into the given ROI of the ImageBuf. The data points to
     /// values specified by `format`, with layout detailed by the stride
-    /// values (in bytes, with AutoStride indicating "contiguous" layout).
-    /// It is up to the caller to ensure that data points to an area of
-    /// memory big enough to account for the ROI. Return true if the
-    /// operation could be completed, otherwise return false.
+    /// values (in bytes, with AutoStride indicating "contiguous" layout). It
+    /// is up to the caller to ensure that data points to an area of memory
+    /// big enough to account for the ROI. If `roi` is set to `ROI::all()`,
+    /// the data buffer is assumed to have the same resolution as the ImageBuf
+    /// itself. Return true if the operation could be completed, otherwise
+    /// return false.
     bool set_pixels(ROI roi, TypeDesc format, const void* data,
                     stride_t xstride = AutoStride,
                     stride_t ystride = AutoStride,

src/python/py_imagebuf.cpp

@@ -13,6 +13,55 @@ namespace PyOpenImageIO {
 
 
 
+static ImageBuf
+ImageBuf_from_buffer(const py::buffer& buffer)
+{
+    ImageBuf ib;
+    const py::buffer_info info = buffer.request();
+    TypeDesc format;
+    if (info.format.size())
+        format = typedesc_from_python_array_code(info.format);
+    if (format == TypeUnknown)
+        return ib;
+    // Strutil::print("IB from {} buffer: dims = {}\n", format, info.ndim);
+    // for (int i = 0; i < info.ndim; ++i)
+    //     Strutil::print("IB from buffer: dim[{}]: size = {}, stride = {}\n", i,
+    //                    info.shape[i], info.strides[i]);
+    if (size_t(info.strides[info.ndim - 1]) != format.size()) {
+        ib.errorfmt(
+            "ImageBuf-from-numpy-array must have contiguous stride within pixels");
+        return ib;
+    }
+
+    if (info.ndim == 3) {
+        // Assume [y][x][c]
+        ImageSpec spec(info.shape[1], info.shape[0], info.shape[2], format);
+        ib.reset(spec, InitializePixels::No);
+        ib.set_pixels(get_roi(spec), format, info.ptr, info.strides[1],
+                      info.strides[0]);
+    } else if (info.ndim == 2) {
+        // Assume [y][x], single channel
+        ImageSpec spec(info.shape[1], info.shape[0], 1, format);
+        ib.reset(spec, InitializePixels::No);
+        ib.set_pixels(get_roi(spec), format, info.ptr, info.strides[1],
+                      info.strides[0]);
+    } else if (info.ndim == 4) {
+        // Assume volume [z][y][x][c]
+        ImageSpec spec(info.shape[2], info.shape[1], info.shape[3], format);
+        spec.depth      = info.shape[0];
+        spec.full_depth = spec.depth;
+        ib.reset(spec, InitializePixels::No);
+        ib.set_pixels(get_roi(spec), format, info.ptr, info.strides[2],
+                      info.strides[1], info.strides[0]);
+    } else {
+        ib.errorfmt(
+            "ImageBuf-from-numpy-array must have 2, 3, or 4 dimensions");
+    }
+    return ib;
+}
+
+
+
 py::tuple
 ImageBuf_getpixel(const ImageBuf& buf, int x, int y, int z = 0,
                   const std::string& wrapname = "black")
@@ -203,6 +252,10 @@ declare_imagebuf(py::module& m)
                  return ImageBuf(name, subimage, miplevel, nullptr, &config);
              }),
              "name"_a, "subimage"_a, "miplevel"_a, "config"_a)
+        .def(py::init([](const py::buffer& buffer) {
+                 return ImageBuf_from_buffer(buffer);
+             }),
+             "buffer"_a)
         .def("clear", &ImageBuf::clear)
         .def(
             "reset",
@@ -224,6 +277,13 @@ declare_imagebuf(py::module& m)
                 self.reset(spec, z);
             },
             "spec"_a, "zero"_a = true)
+        .def(
+            "reset",
+            [](ImageBuf& self, const py::buffer& buffer) {
+                self = ImageBuf_from_buffer(buffer);
+            },
+            "buffer"_a)
+
         .def_property_readonly("initialized",
                                [](const ImageBuf& self) {
                                    return self.initialized();

testsuite/python-imagebuf/ref/out-alt-python3.txt

@@ -14,6 +14,15 @@ Resetting to be a writable 640x480,3 Float:
   alpha channel =  -1
   z channel =  -1
   deep =  False
+Constructing from a bare numpy array:
+  resolution 2x3+0+0
+  untiled
+  4 channels: ('R', 'G', 'B', 'A')
+  format =  float
+  alpha channel =  3
+  z channel =  -1
+  deep =  False
+  pixel (0,1) = 0.3 0 0.8 1
 
 Testing read of ../common/textures/grid.tx:
 subimage: 0  /  1

testsuite/python-imagebuf/ref/out-alt.txt

@@ -14,6 +14,15 @@ Resetting to be a writable 640x480,3 Float:
   alpha channel =  -1
   z channel =  -1
   deep =  False
+Constructing from a bare numpy array:
+  resolution 2x3+0+0
+  untiled
+  4 channels: ('R', 'G', 'B', 'A')
+  format =  float
+  alpha channel =  3
+  z channel =  -1
+  deep =  False
+  pixel (0,1) = 0.3 0 0.8 1
 
 Testing read of ../common/textures/grid.tx:
 subimage: 0  /  1

testsuite/python-imagebuf/ref/out-python3.txt

@@ -14,6 +14,15 @@ Resetting to be a writable 640x480,3 Float:
   alpha channel =  -1
   z channel =  -1
   deep =  False
+Constructing from a bare numpy array:
+  resolution 2x3+0+0
+  untiled
+  4 channels: ('R', 'G', 'B', 'A')
+  format =  float
+  alpha channel =  3
+  z channel =  -1
+  deep =  False
+  pixel (0,1) = 0.3 0 0.8 1
 
 Testing read of ../common/textures/grid.tx:
 subimage: 0  /  1

testsuite/python-imagebuf/ref/out.txt

@@ -14,6 +14,15 @@ Resetting to be a writable 640x480,3 Float:
   alpha channel =  -1
   z channel =  -1
   deep =  False
+Constructing from a bare numpy array:
+  resolution 2x3+0+0
+  untiled
+  4 channels: ('R', 'G', 'B', 'A')
+  format =  float
+  alpha channel =  3
+  z channel =  -1
+  deep =  False
+  pixel (0,1) = 0.3 0 0.8 1
 
 Testing read of ../common/textures/grid.tx:
 subimage: 0  /  1

testsuite/python-imagebuf/src/test_imagebuf.py

@@ -145,6 +145,15 @@ def test_multiimage () :
     print ("Resetting to be a writable 640x480,3 Float:")
     b.reset (oiio.ImageSpec(640,480,3,oiio.FLOAT))
     print_imagespec (b.spec())
+
+    print ("Constructing from a bare numpy array:")
+    b = oiio.ImageBuf(numpy.array([[[0.1,0.0,0.9,1.0], [0.2,0.0,0.7,1.0]],
+                                   [[0.3,0.0,0.8,1.0], [0.4,0.0,0.6,1.0]],
+                                   [[0.5,0.0,0.7,1.0], [0.6,0.0,0.5,1.0]]], dtype="f"))
+    # should be width=2, height=3, channels=4, format=FLOAT
+    print_imagespec (b.spec())
+    print ("  pixel (0,1) = {:.3g} {:.3g} {:.3g} {:.3g}".format(b.getpixel (0,1)[0],
+           b.getpixel (0,1)[1], b.getpixel (0,1)[2], b.getpixel(0,1)[3]))
     print ("")
 
     # Test reading from disk