forked from 170010011/fr
302 lines
11 KiB
Python
302 lines
11 KiB
Python
|
# -*- coding: utf-8 -*-
|
||
|
|
# imageio is distributed under the terms of the (new) BSD License.
|
||
|
|
|
||
|
|
""" BSDF plugin.
|
||
|
|
"""
|
||
|
|
|
||
|
|
import numpy as np
|
||
|
|
|
||
|
|
from .. import formats
|
||
|
|
from ..core import Format
|
||
|
|
|
||
|
|
|
||
|
|
def get_bsdf_serializer(options):
|
||
|
|
from . import _bsdf as bsdf
|
||
|
|
|
||
|
|
class NDArrayExtension(bsdf.Extension):
|
||
|
|
""" Copy of BSDF's NDArrayExtension but deal with lazy blobs.
|
||
|
|
"""
|
||
|
|
|
||
|
|
name = "ndarray"
|
||
|
|
cls = np.ndarray
|
||
|
|
|
||
|
|
def encode(self, s, v):
|
||
|
|
return dict(shape=v.shape, dtype=str(v.dtype), data=v.tobytes())
|
||
|
|
|
||
|
|
def decode(self, s, v):
|
||
|
|
return v # return as dict, because of lazy blobs, decode in Image
|
||
|
|
|
||
|
|
class ImageExtension(bsdf.Extension):
|
||
|
|
""" We implement two extensions that trigger on the Image classes.
|
||
|
|
"""
|
||
|
|
|
||
|
|
def encode(self, s, v):
|
||
|
|
return dict(array=v.array, meta=v.meta)
|
||
|
|
|
||
|
|
def decode(self, s, v):
|
||
|
|
return Image(v["array"], v["meta"])
|
||
|
|
|
||
|
|
class Image2DExtension(ImageExtension):
|
||
|
|
|
||
|
|
name = "image2d"
|
||
|
|
cls = Image2D
|
||
|
|
|
||
|
|
class Image3DExtension(ImageExtension):
|
||
|
|
|
||
|
|
name = "image3d"
|
||
|
|
cls = Image3D
|
||
|
|
|
||
|
|
exts = [NDArrayExtension, Image2DExtension, Image3DExtension]
|
||
|
|
serializer = bsdf.BsdfSerializer(exts, **options)
|
||
|
|
|
||
|
|
return bsdf, serializer
|
||
|
|
|
||
|
|
|
||
|
|
class Image:
|
||
|
|
""" Class in which we wrap the array and meta data. By using an extension
|
||
|
|
we can make BSDF trigger on these classes and thus encode the images.
|
||
|
|
as actual images.
|
||
|
|
"""
|
||
|
|
|
||
|
|
def __init__(self, array, meta):
|
||
|
|
self.array = array
|
||
|
|
self.meta = meta
|
||
|
|
|
||
|
|
def get_array(self):
|
||
|
|
if not isinstance(self.array, np.ndarray):
|
||
|
|
v = self.array
|
||
|
|
blob = v["data"]
|
||
|
|
if not isinstance(blob, bytes): # then it's a lazy bsdf.Blob
|
||
|
|
blob = blob.get_bytes()
|
||
|
|
self.array = np.frombuffer(blob, dtype=v["dtype"])
|
||
|
|
self.array.shape = v["shape"]
|
||
|
|
return self.array
|
||
|
|
|
||
|
|
def get_meta(self):
|
||
|
|
return self.meta
|
||
|
|
|
||
|
|
|
||
|
|
class Image2D(Image):
|
||
|
|
pass
|
||
|
|
|
||
|
|
|
||
|
|
class Image3D(Image):
|
||
|
|
pass
|
||
|
|
|
||
|
|
|
||
|
|
class BsdfFormat(Format):
|
||
|
|
""" The BSDF format enables reading and writing of image data in the
|
||
|
|
BSDF serialization format. This format allows storage of images, volumes,
|
||
|
|
and series thereof. Data can be of any numeric data type, and can
|
||
|
|
optionally be compressed. Each image/volume can have associated
|
||
|
|
meta data, which can consist of any data type supported by BSDF.
|
||
|
|
|
||
|
|
By default, image data is lazily loaded; the actual image data is
|
||
|
|
not read until it is requested. This allows storing multiple images
|
||
|
|
in a single file and still have fast access to individual images.
|
||
|
|
Alternatively, a series of images can be read in streaming mode, reading
|
||
|
|
images as they are read (e.g. from http).
|
||
|
|
|
||
|
|
BSDF is a simple generic binary format. It is easy to extend and there
|
||
|
|
are standard extension definitions for 2D and 3D image data.
|
||
|
|
Read more at http://bsdf.io.
|
||
|
|
|
||
|
|
Parameters for reading
|
||
|
|
----------------------
|
||
|
|
random_access : bool
|
||
|
|
Whether individual images in the file can be read in random order.
|
||
|
|
Defaults to True for normal files, and to False when reading from HTTP.
|
||
|
|
If False, the file is read in "streaming mode", allowing reading
|
||
|
|
files as they are read, but without support for "rewinding".
|
||
|
|
Note that setting this to True when reading from HTTP, the whole file
|
||
|
|
is read upon opening it (since lazy loading is not possible over HTTP).
|
||
|
|
|
||
|
|
Parameters for saving
|
||
|
|
---------------------
|
||
|
|
compression : {0, 1, 2}
|
||
|
|
Use ``0`` or "no" for no compression, ``1`` or "zlib" for Zlib
|
||
|
|
compression (same as zip files and PNG), and ``2`` or "bz2" for Bz2
|
||
|
|
compression (more compact but slower). Default 1 (zlib).
|
||
|
|
Note that some BSDF implementations may not support compression
|
||
|
|
(e.g. JavaScript).
|
||
|
|
|
||
|
|
"""
|
||
|
|
|
||
|
|
def _can_read(self, request):
|
||
|
|
if request.mode[1] in (self.modes + "?"):
|
||
|
|
# if request.extension in self.extensions:
|
||
|
|
# return True
|
||
|
|
if request.firstbytes.startswith(b"BSDF"):
|
||
|
|
return True
|
||
|
|
|
||
|
|
def _can_write(self, request):
|
||
|
|
if request.mode[1] in (self.modes + "?"):
|
||
|
|
if request.extension in self.extensions:
|
||
|
|
return True
|
||
|
|
|
||
|
|
# -- reader
|
||
|
|
|
||
|
|
class Reader(Format.Reader):
|
||
|
|
def _open(self, random_access=None):
|
||
|
|
# Validate - we need a BSDF file consisting of a list of images
|
||
|
|
# The list is typically a stream, but does not have to be.
|
||
|
|
assert self.request.firstbytes[:4] == b"BSDF", "Not a BSDF file"
|
||
|
|
# self.request.firstbytes[5:6] == major and minor version
|
||
|
|
if not (
|
||
|
|
self.request.firstbytes[6:15] == b"M\x07image2D"
|
||
|
|
or self.request.firstbytes[6:15] == b"M\x07image3D"
|
||
|
|
or self.request.firstbytes[6:7] == b"l"
|
||
|
|
):
|
||
|
|
pass # Actually, follow a more duck-type approach ...
|
||
|
|
# raise RuntimeError('BSDF file does not look like an '
|
||
|
|
# 'image container.')
|
||
|
|
# Set options. If we think that seeking is allowed, we lazily load
|
||
|
|
# blobs, and set streaming to False (i.e. the whole file is read,
|
||
|
|
# but we skip over binary blobs), so that we subsequently allow
|
||
|
|
# random access to the images.
|
||
|
|
# If seeking is not allowed (e.g. with a http request), we cannot
|
||
|
|
# lazily load blobs, but we can still load streaming from the web.
|
||
|
|
options = {}
|
||
|
|
if self.request.filename.startswith(("http://", "https://")):
|
||
|
|
ra = False if random_access is None else bool(random_access)
|
||
|
|
options["lazy_blob"] = False # Because we cannot seek now
|
||
|
|
options["load_streaming"] = not ra # Load as a stream?
|
||
|
|
else:
|
||
|
|
ra = True if random_access is None else bool(random_access)
|
||
|
|
options["lazy_blob"] = ra # Don't read data until needed
|
||
|
|
options["load_streaming"] = not ra
|
||
|
|
|
||
|
|
file = self.request.get_file()
|
||
|
|
bsdf, self._serializer = get_bsdf_serializer(options)
|
||
|
|
self._stream = self._serializer.load(file)
|
||
|
|
# Another validation
|
||
|
|
if (
|
||
|
|
isinstance(self._stream, dict)
|
||
|
|
and "meta" in self._stream
|
||
|
|
and "array" in self._stream
|
||
|
|
):
|
||
|
|
self._stream = Image(self._stream["array"], self._stream["meta"])
|
||
|
|
if not isinstance(self._stream, (Image, list, bsdf.ListStream)):
|
||
|
|
raise RuntimeError(
|
||
|
|
"BSDF file does not look seem to have an " "image container."
|
||
|
|
)
|
||
|
|
|
||
|
|
def _close(self):
|
||
|
|
pass
|
||
|
|
|
||
|
|
def _get_length(self):
|
||
|
|
if isinstance(self._stream, Image):
|
||
|
|
return 1
|
||
|
|
elif isinstance(self._stream, list):
|
||
|
|
return len(self._stream)
|
||
|
|
elif self._stream.count < 0:
|
||
|
|
return np.inf
|
||
|
|
return self._stream.count
|
||
|
|
|
||
|
|
def _get_data(self, index):
|
||
|
|
# Validate
|
||
|
|
if index < 0 or index >= self.get_length():
|
||
|
|
raise IndexError(
|
||
|
|
"Image index %i not in [0 %i]." % (index, self.get_length())
|
||
|
|
)
|
||
|
|
# Get Image object
|
||
|
|
if isinstance(self._stream, Image):
|
||
|
|
image_ob = self._stream # singleton
|
||
|
|
elif isinstance(self._stream, list):
|
||
|
|
# Easy when we have random access
|
||
|
|
image_ob = self._stream[index]
|
||
|
|
else:
|
||
|
|
# For streaming, we need to skip over frames
|
||
|
|
if index < self._stream.index:
|
||
|
|
raise IndexError(
|
||
|
|
"BSDF file is being read in streaming "
|
||
|
|
"mode, thus does not allow rewinding."
|
||
|
|
)
|
||
|
|
while index > self._stream.index:
|
||
|
|
self._stream.next()
|
||
|
|
image_ob = self._stream.next() # Can raise StopIteration
|
||
|
|
# Is this an image?
|
||
|
|
if (
|
||
|
|
isinstance(image_ob, dict)
|
||
|
|
and "meta" in image_ob
|
||
|
|
and "array" in image_ob
|
||
|
|
):
|
||
|
|
image_ob = Image(image_ob["array"], image_ob["meta"])
|
||
|
|
if isinstance(image_ob, Image):
|
||
|
|
# Return as array (if we have lazy blobs, they are read now)
|
||
|
|
return image_ob.get_array(), image_ob.get_meta()
|
||
|
|
else:
|
||
|
|
r = repr(image_ob)
|
||
|
|
r = r if len(r) < 200 else r[:197] + "..."
|
||
|
|
raise RuntimeError("BSDF file contains non-image " + r)
|
||
|
|
|
||
|
|
def _get_meta_data(self, index): # pragma: no cover
|
||
|
|
return {} # This format does not support global meta data
|
||
|
|
|
||
|
|
# -- writer
|
||
|
|
|
||
|
|
class Writer(Format.Writer):
|
||
|
|
def _open(self, compression=1):
|
||
|
|
options = {"compression": compression}
|
||
|
|
bsdf, self._serializer = get_bsdf_serializer(options)
|
||
|
|
if self.request.mode[1] in "iv":
|
||
|
|
self._stream = None # Singleton image
|
||
|
|
self._written = False
|
||
|
|
else:
|
||
|
|
# Series (stream) of images
|
||
|
|
file = self.request.get_file()
|
||
|
|
self._stream = bsdf.ListStream()
|
||
|
|
self._serializer.save(file, self._stream)
|
||
|
|
|
||
|
|
def _close(self):
|
||
|
|
# We close the stream here, which will mark the number of written
|
||
|
|
# elements. If we would not close it, the file would be fine, it's
|
||
|
|
# just that upon reading it would not be known how many items are
|
||
|
|
# in there.
|
||
|
|
if self._stream is not None:
|
||
|
|
self._stream.close(False) # False says "keep this a stream"
|
||
|
|
|
||
|
|
def _append_data(self, im, meta):
|
||
|
|
# Determine dimension
|
||
|
|
ndim = None
|
||
|
|
if self.request.mode[1] in "iI":
|
||
|
|
ndim = 2
|
||
|
|
elif self.request.mode[1] in "vV":
|
||
|
|
ndim = 3
|
||
|
|
else:
|
||
|
|
ndim = 3 # Make an educated guess
|
||
|
|
if im.ndim == 2 or (im.ndim == 3 and im.shape[-1] <= 4):
|
||
|
|
ndim = 2
|
||
|
|
# Validate shape
|
||
|
|
assert ndim in (2, 3)
|
||
|
|
if ndim == 2:
|
||
|
|
assert im.ndim == 2 or (im.ndim == 3 and im.shape[-1] <= 4)
|
||
|
|
else:
|
||
|
|
assert im.ndim == 3 or (im.ndim == 4 and im.shape[-1] <= 4)
|
||
|
|
# Wrap data and meta data in our special class that will trigger
|
||
|
|
# the BSDF image2D or image3D extension.
|
||
|
|
if ndim == 2:
|
||
|
|
ob = Image2D(im, meta)
|
||
|
|
else:
|
||
|
|
ob = Image3D(im, meta)
|
||
|
|
# Write directly or to stream
|
||
|
|
if self._stream is None:
|
||
|
|
assert not self._written, "Cannot write singleton image twice"
|
||
|
|
self._written = True
|
||
|
|
file = self.request.get_file()
|
||
|
|
self._serializer.save(file, ob)
|
||
|
|
else:
|
||
|
|
self._stream.append(ob)
|
||
|
|
|
||
|
|
def set_meta_data(self, meta): # pragma: no cover
|
||
|
|
raise RuntimeError("The BSDF format only supports " "per-image meta data.")
|
||
|
|
|
||
|
|
|
||
|
|
format = BsdfFormat(
|
||
|
|
"bsdf", # short name
|
||
|
|
"Format based on the Binary Structured Data Format",
|
||
|
|
".bsdf",
|
||
|
|
"iIvV",
|
||
|
|
)
|
||
|
|
formats.add_format(format)
|