86 lines
2.6 KiB
Python
86 lines
2.6 KiB
Python
|
# -*- coding: utf-8 -*-
|
||
|
|
# imageio is distributed under the terms of the (new) BSD License.
|
||
|
|
|
||
|
|
"""Read/Write NPZ files.
|
||
|
|
|
||
|
|
Backend: `Numpy <https://numpy.org/doc/stable/reference/generated/numpy.savez.html>`_
|
||
|
|
|
||
|
|
NPZ is a file format by numpy that provides storage of array data using gzip
|
||
|
|
compression. This imageio plugin supports data of any shape, and also supports
|
||
|
|
multiple images per file. However, the npz format does not provide streaming;
|
||
|
|
all data is read/written at once. Further, there is no support for meta data.
|
||
|
|
|
||
|
|
See the BSDF format for a similar (but more fully featured) format.
|
||
|
|
|
||
|
|
Parameters
|
||
|
|
----------
|
||
|
|
None
|
||
|
|
|
||
|
|
Notes
|
||
|
|
-----
|
||
|
|
This format is not available on Pypy.
|
||
|
|
|
||
|
|
"""
|
||
|
|
|
||
|
|
import numpy as np
|
||
|
|
|
||
|
|
from ..core import Format
|
||
|
|
|
||
|
|
|
||
|
|
class NpzFormat(Format):
|
||
|
|
"""See :mod:`imageio.plugins.npz`"""
|
||
|
|
|
||
|
|
def _can_read(self, request):
|
||
|
|
# We support any kind of image data
|
||
|
|
return request.extension in self.extensions
|
||
|
|
|
||
|
|
def _can_write(self, request):
|
||
|
|
# We support any kind of image data
|
||
|
|
return request.extension in self.extensions
|
||
|
|
|
||
|
|
# -- reader
|
||
|
|
|
||
|
|
class Reader(Format.Reader):
|
||
|
|
def _open(self):
|
||
|
|
# Load npz file, which provides another file like object
|
||
|
|
self._npz = np.load(self.request.get_file())
|
||
|
|
assert isinstance(self._npz, np.lib.npyio.NpzFile)
|
||
|
|
# Get list of names, ordered by name, but smarter
|
||
|
|
self._names = sorted(self._npz.files, key=lambda x: x.split("_")[-1])
|
||
|
|
|
||
|
|
def _close(self):
|
||
|
|
self._npz.close()
|
||
|
|
|
||
|
|
def _get_length(self):
|
||
|
|
return len(self._names)
|
||
|
|
|
||
|
|
def _get_data(self, index):
|
||
|
|
# Get data
|
||
|
|
if index < 0 or index >= len(self._names):
|
||
|
|
raise IndexError("Index out of range while reading from nzp")
|
||
|
|
im = self._npz[self._names[index]]
|
||
|
|
# Return array and empty meta data
|
||
|
|
return im, {}
|
||
|
|
|
||
|
|
def _get_meta_data(self, index):
|
||
|
|
# Get the meta data for the given index
|
||
|
|
raise RuntimeError("The npz format does not support meta data.")
|
||
|
|
|
||
|
|
# -- writer
|
||
|
|
|
||
|
|
class Writer(Format.Writer):
|
||
|
|
def _open(self):
|
||
|
|
# Npz is not such a great format. We cannot stream to the file.
|
||
|
|
# So we remember all images and write them to file at the end.
|
||
|
|
self._images = []
|
||
|
|
|
||
|
|
def _close(self):
|
||
|
|
# Write everything
|
||
|
|
np.savez_compressed(self.request.get_file(), *self._images)
|
||
|
|
|
||
|
|
def _append_data(self, im, meta):
|
||
|
|
self._images.append(im) # discart meta data
|
||
|
|
|
||
|
|
def set_meta_data(self, meta):
|
||
|
|
raise RuntimeError("The npz format does not support meta data.")
|