def write_png(data, origin='<STR_LIT>', colormap=None):
if colormap is None:<EOL><INDENT>def colormap(x):<EOL><INDENT>return (x, x, x, <NUM_LIT:1>)<EOL><DEDENT><DEDENT>arr = np.atleast_3d(data)<EOL>height, width, nblayers = arr.shape<EOL>if nblayers not in [<NUM_LIT:1>, <NUM_LIT:3>, <NUM_LIT:4>]:<EOL><INDENT>raise ValueError('<STR_LIT>'<EOL>'<STR_LIT>')<EOL><DEDENT>assert arr.shape == (height, width, nblayers)<EOL>if nblayers == <NUM_LIT:1>:<EOL><INDENT>arr = np.array(list(map(colormap, arr.ravel())))<EOL>nblayers = arr.shape[<NUM_LIT:1>]<EOL>if nblayers not in [<NUM_LIT:3>, <NUM_LIT:4>]:<EOL><INDENT>raise ValueError('<STR_LIT>'<EOL>'<STR_LIT>')<EOL><DEDENT>arr = arr.reshape((height, width, nblayers))<EOL><DEDENT>assert arr.shape == (height, width, nblayers)<EOL>if nblayers == <NUM_LIT:3>:<EOL><INDENT>arr = np.concatenate((arr, np.ones((height, width, <NUM_LIT:1>))), axis=<NUM_LIT:2>)<EOL>nblayers = <NUM_LIT:4><EOL><DEDENT>assert arr.shape == (height, width, nblayers)<EOL>assert nblayers == <NUM_LIT:4><EOL>if arr.dtype != '<STR_LIT>':<EOL><INDENT>with np.errstate(divide='<STR_LIT:ignore>', invalid='<STR_LIT:ignore>'):<EOL><INDENT>arr = arr * <NUM_LIT>/arr.max(axis=(<NUM_LIT:0>, <NUM_LIT:1>)).reshape((<NUM_LIT:1>, <NUM_LIT:1>, <NUM_LIT:4>))<EOL>arr[~np.isfinite(arr)] = <NUM_LIT:0><EOL><DEDENT>arr = arr.astype('<STR_LIT>')<EOL><DEDENT>if origin == '<STR_LIT>':<EOL><INDENT>arr = arr[::-<NUM_LIT:1>, :, :]<EOL><DEDENT>raw_data = b'<STR_LIT>'.join([b'<STR_LIT:\x00>' + arr[i, :, :].tobytes()<EOL>for i in range(height)])<EOL>def png_pack(png_tag, data):<EOL><INDENT>chunk_head = png_tag + data<EOL>return (struct.pack('<STR_LIT>', len(data)) +<EOL>chunk_head +<EOL>struct.pack('<STR_LIT>', <NUM_LIT> & zlib.crc32(chunk_head)))<EOL><DEDENT>return b'<STR_LIT>'.join([<EOL>b'<STR_LIT>',<EOL>png_pack(b'<STR_LIT>', struct.pack('<STR_LIT>', width, height, <NUM_LIT:8>, <NUM_LIT:6>, <NUM_LIT:0>, <NUM_LIT:0>, <NUM_LIT:0>)),<EOL>png_pack(b'<STR_LIT>', zlib.compress(raw_data, <NUM_LIT:9>)),<EOL>png_pack(b'<STR_LIT>', b'<STR_LIT>')])<EOL>
Transform an array of data into a PNG string.
This can be written to disk using binary I/O, or encoded using base64
for an inline PNG like this:

>>> png_str = write_png(array)
>>> "data:image/png;base64,"+png_str.encode('base64')

Inspired from
https://stackoverflow.com/questions/902761/saving-a-numpy-array-as-an-image

Parameters
----------
data: numpy array or equivalent list-like object.
     Must be NxM (mono), NxMx3 (RGB) or NxMx4 (RGBA)

origin : ['upper' | 'lower'], optional, default 'upper'
    Place the [0,0] index of the array in the upper left or lower left
    corner of the axes.

colormap : callable, used only for `mono` image.
    Function of the form [x -> (r,g,b)] or [x -> (r,g,b,a)]
    for transforming a mono image into RGB.
    It must output iterables of length 3 or 4, with values between
    0. and 1.  Hint: you can use colormaps from `matplotlib.cm`.

Returns
-------
PNG formatted byte string