halide.imageio needs to support arbitrary bufferviews (halide#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)