halide.imageio needs to support arbitrary bufferviews (#7137)

* halide.imageio needs to support arbitrary bufferviews

As written, the helper code assumed that everything passed in was a numpy array of some sort; this meant that passing hl.Buffer didn't work. Restructured so that we only assume that the objects passed in satisfies the Python buffer protocol, so this should now work very generically.

* Update imageio.py

* More fixes

python_bindings/src/halide/imageio.py

@@ -2,72 +2,59 @@
     import imageio.v2 as imageio
 except:
     import imageio
-import numpy
+import numpy as np
 
 
 def is_interleaved(im):
-    """If the given ndarray is 3-dimensional and appears to have an interleaved
-       layout, return True. Otherwise, return False."""
+    """If the given buffer is 3-dimensional and appears to have an interleaved
+    layout, return True. Otherwise, return False."""
 
     # Assume that 'interleaved' will only apply to channels <= 4
-    return im.ndim == 3 and im.strides[2] == 1 and im.shape[2] in [1, 2, 3, 4]
+    mv = memoryview(im)
+    return mv.ndim == 3 and mv.strides[2] == 1 and mv.shape[2] in [1, 2, 3, 4]
 
 
-def _as_interleaved(im):
-    """If the given ndarray is 3-dimensional and appears to be planar layout,
-       return a view that is in interleaved form, leaving the input unchanged.
-       Otherwise, return the image ndarray unchanged.
-       Note that this call must be used with care, as the returnee may or may
-       not be a copy."""
-    if im.ndim == 3 and not is_interleaved(im):
-        return numpy.moveaxis(im, 0, 2)
+def copy_to_interleaved(im):
+    """If the given buffer is 3-dimensional and appears to be planar
+    layout, return a copy that is in interleaved form. Otherwise, return
+    an unchanged copy of the input. Note that this call will always return
+    a copy, leaving the input unchanged."""
+    mv = memoryview(im)
+    if mv.ndim == 3 and not is_interleaved(mv):
+        # We are presumably planar, in (c, y, x) order; we need (y, x, c) order
+        mv = np.moveaxis(mv, 0, 2)
+        mv = np.copy(mv, order="F")
+        return mv
     else:
         return im
 
 
-def _as_planar(im):
-    """If the given ndarray is 3-dimensional and appears to be interleaved
-       layout, return a view that is in planar form, leaving the input
-       unchanged. Otherwise, return the image ndarray unchanged.
-       Note that this call must be used with care, as the returnee may or may
-       not be a copy."""
-    if is_interleaved(im):
-        return numpy.moveaxis(im, 2, 0)
+def copy_to_planar(im):
+    """If the given buffer is 3-dimensional and appears to be interleaved
+    layout, return a copy that is in planar form. Otherwise, return
+    an unchanged copy of the input. Note that this call will always return
+    a copy, leaving the input unchanged."""
+    mv = memoryview(im)
+    if is_interleaved(mv):
+        # Interleaved will be in (y, x, c) order; we want (c, y, x) order
+        # (which hl.Buffer will reverse into x, y, c order)
+        mv = np.moveaxis(mv, 2, 0)
+        mv = np.copy(mv, order="C")
+        return mv
     else:
         return im
 
 
-def copy_to_interleaved(im):
-    """If the given ndarray is 3-dimensional and appears to be planar
-       layout, return a copy that is in interleaved form. Otherwise, return
-       an unchanged copy of the input. Note that this call will always return
-       a copy, leaving the input unchanged."""
-    return _as_interleaved(im).copy()
-
-
-def copy_to_planar(im):
-    """If the given ndarray is 3-dimensional and appears to be interleaved
-       layout, return a copy that is in planar form. Otherwise, return
-       an unchanged copy of the input. Note that this call will always return
-       a copy, leaving the input unchanged."""
-    return _as_planar(im).copy()
-
-
 def imread(uri, format=None, **kwargs):
     """halide.imageio.imread is a thin wrapper around imagio.imread,
-       except that for 3-dimensional images that appear to be interleaved,
-       the result is converted to a planar layout before returning."""
+    except that for 3-dimensional images that appear to be interleaved,
+    the result is converted to a planar layout before returning."""
     return copy_to_planar(imageio.imread(uri, format, **kwargs))
 
 
 def imwrite(uri, im, format=None, **kwargs):
     """halide.imageio.imwrite is a thin wrapper around imagio.imwrite,
-       except that for 3-dimensional images that appear to be planar,
-       a temporary interleaved copy of the input is made, which is used for
-       writing."""
-
-    # We can use _as_interleaved() here to save a possible copy; since the
-    # caller will never see the possibly-a-copy value, there should be no
-    # risk of possibly-different behavior between cases that need converting
-    # and cases that don't.
-    imageio.imwrite(uri, _as_interleaved(im), format, **kwargs)
+    except that for 3-dimensional images that appear to be planar,
+    a temporary interleaved copy of the input is made, which is used for
+    writing."""
+    imageio.imwrite(uri, copy_to_interleaved(im), format, **kwargs)