# Copyright 2020 Timothy M. Shead
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Helpers for implementing Imagecat I/O.
"""
import bz2
import collections
import json
import logging
import os
import pickle
import sys
import numpy
import imagecat.color
import imagecat.data
import imagecat.optional
import imagecat.require
Imath = imagecat.optional.module("Imath")
OpenEXR = imagecat.optional.module("OpenEXR")
PIL = imagecat.optional.module("PIL.Image")
log = logging.getLogger(__name__)
#####################################################################################33
# Loaders
[docs]@imagecat.require.loaded_module(["Imath", "OpenEXR"])
def openexr_loader(task, path, layers):
"""Image loader plugin for OpenEXR (.exr) files.
Implemented using https://www.excamera.com/sphinx/articles-openexr.html
Use :func:`imagecat.operator.load` to load images in an Imagecat workflow.
"""
extension = os.path.splitext(path)[1].lower()
if extension != ".exr":
return None
if layers != "*":
raise NotImplementedError("Layer matching not implemented.")
reader = OpenEXR.InputFile(path)
header = reader.header()
width = header["dataWindow"].max.x - header["dataWindow"].min.x + 1
height = header["dataWindow"].max.y - header["dataWindow"].min.y + 1
# Cleanup the OpenEXR header to use as the image metadata.
class CustomEncoder(json.JSONEncoder):
def default(self, o):
if isinstance(o, bytes):
return o.decode("UTF-8")
elif isinstance(o, Imath.Channel):
return (o.type, o.xSampling, o.ySampling)
elif isinstance(o, (Imath.Compression, Imath.LineOrder, Imath.PixelType)):
return str(o)
elif isinstance(o, (Imath.V2f, Imath.V2i, Imath.point)):
return (o.x, o.y)
elif isinstance(o, (Imath.Box, Imath.Box2f, Imath.Box2i)):
return (o.min.x, o.min.y, o.max.x, o.max.y)
return super().default(o)
metadata = json.dumps(header, cls=CustomEncoder)
metadata = json.loads(metadata)
# Load each OpenEXR channel into a separate Imagecat Layer.
layers = {}
for name, dtype in header["channels"].items():
if dtype.type.v == Imath.PixelType.HALF:
data = numpy.frombuffer(reader.channel(name), dtype=numpy.float16).reshape((height, width, 1))
elif dtype.type.v == Imath.PixelType.FLOAT:
data = numpy.frombuffer(reader.channel(name), dtype=numpy.float32).reshape((height, width, 1))
elif dtype.type.v == Imath.PixelType.INT:
data = numpy.frombuffer(reader.channel(name), dtype=numpy.int32).reshape((height, width, 1))
layers[name] = imagecat.data.Layer(data=data, role=imagecat.data.Role.NONE)
return imagecat.data.Image(layers=layers, metadata=metadata)
[docs]def pickle_loader(task, path, layers):
"""Image loader plugin for Imagecat Pickle (.icp) files.
The .icp format serializes an Imagecat image as a gzip2-compressed Python
pickle object. It is primarily used in testing, and is **not** recommended
for general use.
Use :func:`imagecat.operator.load` to load images in an Imagecat workflow.
"""
extension = os.path.splitext(path)[1].lower()
if extension != ".icp":
return None
if layers != "*":
raise NotImplementedError("Layer matching not implemented.")
with bz2.open(path, "rb") as stream:
image = pickle.load(stream)
if not isinstance(image, imagecat.data.Image):
raise RuntimeError("Not an Imagecat Pickle (*.icp) file.") # pragma: no cover
return image
[docs]@imagecat.require.loaded_module("PIL.Image")
def pil_loader(task, path, layers):
"""Image loader plugin that uses Pillow for file I/O.
Loads any file format supported by Pillow, https://pillow.readthedocs.io.
Use :func:`imagecat.operator.load` to load images in an Imagecat workflow.
"""
if layers != "*":
raise NotImplementedError("Layer matching not implemented.")
pil_image = PIL.Image.open(path)
log.debug(pil_image.info)
image = imagecat.data.Image()
if pil_image.mode == "L":
image.layers["Y"] = imagecat.data.Layer(data=numpy.array(pil_image, dtype=numpy.float16)[:,:,None] / 255.0, role=imagecat.data.Role.LUMINANCE)
if pil_image.mode == "RGB":
image.layers["C"] = imagecat.data.Layer(data=imagecat.color.srgb_to_linear(numpy.array(pil_image, dtype=numpy.float16) / 255.0), role=imagecat.data.Role.RGB)
if pil_image.mode == "RGBA":
image.layers["C"] = imagecat.data.Layer(data=imagecat.color.srgb_to_linear(numpy.array(pil_image, dtype=numpy.float16)[:,:,0:3] / 255.0), role=imagecat.data.Role.RGB)
image.layers["A"] = imagecat.data.Layer(data=numpy.array(pil_image, dtype=numpy.float16)[:,:,3:4] / 255.0)
return image
#####################################################################################33
# Savers
[docs]@imagecat.require.loaded_module(["Imath", "OpenEXR"])
def openexr_saver(task, image, layers, path):
"""Image saver plugin for OpenEXR (.exr) files.
Implemented using https://www.excamera.com/sphinx/articles-openexr.html
Use :func:`imagecat.operator.save` to save images in an Imagecat workflow.
"""
extension = os.path.splitext(path)[1].lower()
if extension != ".exr":
return False
exr_channels = {}
exr_pixels = {}
for layer_name in layers:
layer = image.layers[layer_name]
dtype = layer.data.dtype
shape = layer.data.shape
for index, channel in enumerate(imagecat.data.channels(layer.role)):
channel_name = f"{layer_name}.{channel}" if channel else f"{layer_name}"
if dtype == numpy.float16:
exr_channels[channel_name] = Imath.Channel(Imath.PixelType(Imath.PixelType.HALF))
elif dtype == numpy.float32:
exr_channels[channel_name] = Imath.Channel(Imath.PixelType(Imath.PixelType.FLOAT))
elif dtype == numpy.int32:
exr_channels[channel_name] = Imath.Channel(Imath.PixelType(Imath.PixelType.INT))
else:
raise ValueError(f"Unsupported dtype: {dtype}")
exr_pixels[channel_name] = layer.data[:,:,index].tobytes()
exr_header = OpenEXR.Header(shape[1], shape[0])
exr_header["channels"] = exr_channels
exr_writer = OpenEXR.OutputFile(path, exr_header)
exr_writer.writePixels(exr_pixels)
return True
[docs]def pickle_saver(task, image, layers, path):
"""Image saver plugin for Imagecat Pickle (.icp) files.
The .icp format serializes an Imagecat image as a gzip2-compressed Python
pickle object. It is primarily used in testing, and is **not** recommended
for general use.
Use :func:`imagecat.operator.save` to save images in an Imagecat workflow.
"""
extension = os.path.splitext(path)[1].lower()
if extension != ".icp":
return False
with bz2.open(path, "wb") as stream:
pickle.dump(image, stream)
return True
[docs]@imagecat.require.loaded_module("PIL.Image")
def pil_saver(task, image, layers, path):
"""Image saver plugin that uses Pillow for file I/O.
Saves any file format supported by Pillow, https://pillow.readthedocs.io.
Use :func:`imagecat.operator.save` to save images in an Imagecat workflow.
"""
# PIL can only save single-layer images.
if len(layers) != 1:
return False
# See if PIL can save this file format.
base, extension = os.path.splitext(path)
if extension not in PIL.Image.registered_extensions().keys():
return False
pil_image = imagecat.data.to_pil(image.layers[layers[0]])
pil_image.save(path)
return True
loaders = [
pickle_loader,
]
"""List of available loader plugins. In-house plugins may be prepended to this list for use with :func:`imagecat.operator.load`."""
if "OpenEXR" in sys.modules:
loaders.append(openexr_loader)
if "PIL.Image" in sys.modules:
loaders.append(pil_loader)
savers = [
pickle_saver,
]
"""List of available saver plugins. In-house plugins may be prepended to this list for use with :func:`imagecat.operator.save`."""
if "OpenEXR" in sys.modules:
savers.append(openexr_saver)
if "PIL.Image" in sys.modules:
savers.append(pil_saver)