npz.py 2.6 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485
  1. # -*- coding: utf-8 -*-
  2. # imageio is distributed under the terms of the (new) BSD License.
  3. """Read/Write NPZ files.
  4. Backend: `Numpy <https://numpy.org/doc/stable/reference/generated/numpy.savez.html>`_
  5. NPZ is a file format by numpy that provides storage of array data using gzip
  6. compression. This imageio plugin supports data of any shape, and also supports
  7. multiple images per file. However, the npz format does not provide streaming;
  8. all data is read/written at once. Further, there is no support for meta data.
  9. See the BSDF format for a similar (but more fully featured) format.
  10. Parameters
  11. ----------
  12. None
  13. Notes
  14. -----
  15. This format is not available on Pypy.
  16. """
  17. import numpy as np
  18. from ..core import Format
  19. class NpzFormat(Format):
  20. """See :mod:`imageio.plugins.npz`"""
  21. def _can_read(self, request):
  22. # We support any kind of image data
  23. return request.extension in self.extensions
  24. def _can_write(self, request):
  25. # We support any kind of image data
  26. return request.extension in self.extensions
  27. # -- reader
  28. class Reader(Format.Reader):
  29. def _open(self):
  30. # Load npz file, which provides another file like object
  31. self._npz = np.load(self.request.get_file())
  32. assert isinstance(self._npz, np.lib.npyio.NpzFile)
  33. # Get list of names, ordered by name, but smarter
  34. self._names = sorted(self._npz.files, key=lambda x: x.split("_")[-1])
  35. def _close(self):
  36. self._npz.close()
  37. def _get_length(self):
  38. return len(self._names)
  39. def _get_data(self, index):
  40. # Get data
  41. if index < 0 or index >= len(self._names):
  42. raise IndexError("Index out of range while reading from nzp")
  43. im = self._npz[self._names[index]]
  44. # Return array and empty meta data
  45. return im, {}
  46. def _get_meta_data(self, index):
  47. # Get the meta data for the given index
  48. raise RuntimeError("The npz format does not support meta data.")
  49. # -- writer
  50. class Writer(Format.Writer):
  51. def _open(self):
  52. # Npz is not such a great format. We cannot stream to the file.
  53. # So we remember all images and write them to file at the end.
  54. self._images = []
  55. def _close(self):
  56. # Write everything
  57. np.savez_compressed(self.request.get_file(), *self._images)
  58. def _append_data(self, im, meta):
  59. self._images.append(im) # discart meta data
  60. def set_meta_data(self, meta):
  61. raise RuntimeError("The npz format does not support meta data.")