# -*- mode: python; coding: utf-8 -*-
# Copyright 2020 the AAS WorldWide Telescope project
# Licensed under the MIT License.
"""
Low-level loading and handling of images.
Here, images are defined as 2D data buffers stored in memory. Images might be
small, if dealing with an individual tile, or extremely large, if loading up a
large study for tiling.
"""
from __future__ import absolute_import, division, print_function
__all__ = '''
Image
ImageLoader
ImageMode
'''.split()
from enum import Enum
from PIL import Image as pil_image
import numpy as np
import sys
[docs]class ImageMode(Enum):
"""
Allowed image "modes", describing their pixel data formats.
These align with PIL modes when possible, but we expect to need to support
modes that aren't present in PIL (namely, float64). There are also various
obscure PIL modes that we do not support.
"""
RGB = 'RGB'
RGBA = 'RGBA'
F32 = 'F'
[docs] def get_default_save_extension(self):
"""
Get the file extension to be used in this mode's "default" save format.
Returns
-------
The extension, without a period; either "png" or "npy"
"""
# For RGB, could use JPG? But would complexify lots of things downstream.
if self in (ImageMode.RGB, ImageMode.RGBA):
return 'png'
elif self == ImageMode.F32:
return 'npy'
else:
raise Exception('unhandled mode in get_default_save_extension')
[docs] def make_maskable_buffer(self, buf_height, buf_width):
"""
Return a new, uninitialized buffer of the specified shape, with a mode
compatible with this one but able to accept undefined values.
Parameters
----------
buf_height : int
The height of the new buffer
buf_width : int
The width of the new buffer
Returns
-------
An uninitialized :class:`Image` instance.
Notes
-----
"Maskable" means that the buffer can accommodate undefined values.
If the image is RGB or RGBA, that means that the buffer will have an
alpha channel. If the image is scientific, that means that the buffer
will be able to accept NaNs.
"""
mask_mode = self
if self in (ImageMode.RGB, ImageMode.RGBA):
arr = np.empty((buf_height, buf_width, 4), dtype=np.uint8)
mask_mode = ImageMode.RGBA
elif self == ImageMode.F32:
arr = np.empty((buf_height, buf_width), dtype=np.float32)
else:
raise Exception('unhandled mode in make_maskable_buffer()')
return Image.from_array(mask_mode, arr)
[docs] def try_as_pil(self):
"""
Attempt to convert this mode into a PIL image mode string.
Returns
-------
A PIL image mode string, or None if there is no exact counterpart.
"""
return self.value
[docs]class ImageLoader(object):
"""
A class defining how to load an image.
This is implemented as its own class since there can be some options
involved, and we want to provide a centralized place for handling them all.
TODO: support FITS, Numpy, etc.
"""
black_to_transparent = False
colorspace_processing = 'srgb'
desired_mode = None
psd_single_layer = None
[docs] @classmethod
def add_arguments(cls, parser):
"""
Add standard image-loading options to an argparse parser object.
Parameters
----------
parser : :class:`argparse.ArgumentParser`
The argument parser to modify
Returns
-------
The :class:`ImageLoader` class (for chainability).
Notes
-----
If you are writing a command-line interface that takes a single image as
an input, use this function to wire in to standardized image-loading
infrastructure and options.
"""
parser.add_argument(
'--black-to-transparent',
action = 'store_true',
help = 'Convert full black colors to be transparent',
)
parser.add_argument(
'--colorspace-processing',
metavar = 'MODE',
default = 'srgb',
help = 'What kind of RGB colorspace processing to perform (default: %(default)s; choices: %(choices)s)',
choices = ['srgb', 'none'],
)
# not exposing desired_mode -- shouldn't be something the for the user to deal with
parser.add_argument(
'--psd-single-layer',
type = int,
metavar = 'NUMBER',
help = 'If loading a Photoshop image, the (0-based) layer number to load -- saves memory',
)
return cls
[docs] @classmethod
def create_from_args(cls, settings):
"""
Process standard image-loading options to create an :class:`ImageLoader`.
Parameters
----------
settings : :class:`argparse.Namespace`
Settings from processing command-line arguments
Returns
-------
A new :class:`ImageLoader` initialized with the settings.
"""
loader = cls()
loader.black_to_transparent = settings.black_to_transparent
loader.colorspace_processing = settings.colorspace_processing
loader.psd_single_layer = settings.psd_single_layer
return loader
[docs] def load_pil(self, pil_img):
"""
Load an already opened PIL image.
Parameters
----------
pil_img : :class:`PIL.Image.Image`
The image.
Returns
-------
A new :class:`Image`.
Notes
-----
This function should be used instead of :meth:`Image.from_pil` because
may postprocess the image in various ways, depending on the loader
configuration.
"""
# If 8-bit grayscale, convert to RGB. Added for
# https://www.flickr.com/photos/10795027@N08/43023455582 . That file
# also comes with an ICC profile so we do the conversion before
# colorspace processing. TODO: if/when we do FITS to RGB conversion, we
# should probably use that codepath rather than this manual hack.
if pil_img.mode == 'L':
pil_img = pil_img.convert('RGBA')
# Convert pure black to transparent -- make sure to do this before any
# colorspace processing.
#
# As far as I can tell, PIL has no good way to modify an image on a
# pixel-by-pixel level in-place, which is really annoying. For now I'm
# doing this processing in Numpy space, but for an image that almost
# fills up memory we might have to use a different approach since this
# one will involve holding two buffers at once.
if self.black_to_transparent:
if pil_img.mode != 'RGBA':
pil_img = pil_img.convert('RGBA')
a = np.asarray(pil_img)
a = a.copy() # read-only buffer => writeable
for i in range(a.shape[0]):
nonblack = (a[i,...,0] > 0)
np.logical_or(nonblack, a[i,...,1] > 0, out=nonblack)
np.logical_or(nonblack, a[i,...,2] > 0, out=nonblack)
a[i,...,3] *= nonblack
# This is my attempt to preserve the image metadata and other
# attributes, swapping out the pixel data only. There is probably
# a better way to do this
new_img = pil_image.fromarray(a, mode=pil_img.mode)
pil_img.im = new_img.im
del a, new_img
# Make sure that we end up in the right color space. From experience, some
# EPO images have funky colorspaces and we need to convert to sRGB to get
# the tiled versions to appear correctly.
if self.colorspace_processing != 'none' and 'icc_profile' in pil_img.info:
assert self.colorspace_processing == 'srgb' # more modes, one day?
try:
from PIL import ImageCms
except ImportError:
print('''warning: colorspace processing requested, but no `ImageCms` module found in PIL.
Your installation of PIL probably does not have colorspace support.
Colors will not be transformed to sRGB and therefore may not appear as intended.
Compare toasty's output to your source image and decide if this is acceptable to you.
Consider a different setting of the `--colorspace-processing` argument to avoid this warning.''',
file=sys.stderr)
else:
from io import BytesIO
in_prof = ImageCms.getOpenProfile(BytesIO(pil_img.info['icc_profile']))
out_prof = ImageCms.createProfile('sRGB')
xform = ImageCms.buildTransform(in_prof, out_prof, pil_img.mode, pil_img.mode)
ImageCms.applyTransform(pil_img, xform, inPlace=True)
# Make sure that we end up with the right mode, if requested.
if self.desired_mode is not None:
desired_pil = self.desired_mode.try_as_pil()
if desired_pil is None:
raise Exception('cannot convert PIL image to desired mode {}'.format(self.desired_mode))
if pil_img.mode != desired_pil:
pil_img = pil_img.convert(desired_pil)
return Image.from_pil(pil_img)
[docs] def load_stream(self, stream):
"""
Load an image into memory from a file-like stream.
Parameters
----------
stream : file-like
The data to load. Reads should yield bytes.
Returns
-------
A new :class:`Image`.
"""
# TODO: one day, we'll support FITS files and whatnot and we'll have a
# mode where we get a Numpy array but not a PIL image. For now, just
# pass it off to PIL and hope for the best.
# Prevent PIL decompression-bomb aborts. Not thread-safe, of course.
old_max = pil_image.MAX_IMAGE_PIXELS
try:
pil_image.MAX_IMAGE_PIXELS = None
pilimg = pil_image.open(stream)
finally:
pil_image.MAX_IMAGE_PIXELS = old_max
# Now pass it off to generic PIL handling ...
return self.load_pil(pilimg)
[docs] def load_path(self, path):
"""
Load an image into memory from a filesystem path.
Parameters
----------
path : str
The filesystem path to load.
Returns
-------
A new :class:`Image`.
"""
# Special handling for Numpy arrays. TODO: it would be better to sniff
# filetypes instead of just looking at extensions. But, lazy.
if path.endswith('.npy'):
arr = np.load(path)
return Image.from_array(ImageMode.F32, arr.astype(np.float32))
# Special handling for Photoshop files, used for some very large mosaics
# with transparency (e.g. the PHAT M31/M33 images).
if path.endswith('.psd') or path.endswith('.psb'):
try:
from psd_tools import PSDImage
except ImportError:
pass
else:
psd = PSDImage.open(path)
# If the Photoshop image is a single layer, we can save a lot of
# memory by not using the composite() function. This has helped
# me process very large Photoshop files.
if self.psd_single_layer is not None:
pilimg = psd[self.psd_single_layer].topil()
else:
pilimg = psd.composite()
return self.load_pil(pilimg)
# Special handling for OpenEXR files, used for large images with high
# dynamic range.
if path.endswith('.exr'):
from .openexr import load_openexr
img = load_openexr(path)
return Image.from_array(ImageMode.RGB, img)
# (One day, maybe we'll do more kinds of sniffing.) No special handling
# came into play; just open the file and auto-detect.
with open(path, 'rb') as f:
return self.load_stream(f)
[docs]class Image(object):
"""A 2D data array stored in memory.
This class primarily exists to help us abstract between the cases where we
have "bitmap" RGB(A) images and "science" floating-point images.
"""
_mode = None
_pil = None
_array = None
[docs] @classmethod
def from_pil(cls, pil_img):
"""
Create a new Image from a PIL image.
Parameters
----------
pil_img : :class:`PIL.Image.Image`
The source image.
Returns
-------
A new :class:`Image` wrapping the PIL image.
"""
# Make sure that the image data are actually loaded from disk. Pillow
# lazy-loads such that sometimes `np.asarray(img)` ends up failing
# mysteriously if the image is loaded from a file handle that is closed
# promptly.
pil_img.load()
inst = cls()
inst._pil = pil_img
try:
inst._mode = ImageMode(pil_img.mode)
except ValueError:
raise Exception('image mode {} is not supported'.format(pil_img.mode))
return inst
[docs] @classmethod
def from_array(cls, mode, array):
"""Create a new Image from an array-like data variable.
Parameters
----------
mode : :class:`ImageMode`
The image mode.
array : array-like object
The source data.
Returns
-------
A new :class:`Image` wrapping the data.
Notes
-----
The array will be converted to be at least two-dimensional. The data
array shape should match the requirements of the mode.
"""
array = np.atleast_2d(array)
array_ok = False
if mode == ImageMode.F32:
array_ok = (array.ndim == 2 and array.dtype == np.dtype(np.float32))
elif mode == ImageMode.RGB:
array_ok = (array.ndim == 3 and array.shape[2] == 3 and array.dtype == np.dtype(np.uint8))
elif mode == ImageMode.RGBA:
array_ok = (array.ndim == 3 and array.shape[2] == 4 and array.dtype == np.dtype(np.uint8))
else:
raise ValueError('unhandled image mode {} in from_array()'.format(mode))
if not array_ok:
raise ValueError('expected array compatible with mode {}, but got shape/dtype = {}/{}'
.format(mode, array.shape, array.dtype))
inst = cls()
inst._mode = mode
inst._array = array
return inst
[docs] def asarray(self):
"""Obtain the image data as a Numpy array.
Returns
-------
If the image is an RGB(A) bitmap, the array will have shape ``(height, width, planes)``
and a dtype of ``uint8``, where ``planes`` is either 3
or 4 depending on whether the image has an alpha channel. If the image
is science data, it will have shape ``(height, width)`` and have a
floating-point dtype.
"""
# NOTE: it turns out that np.asarray() on a PIL image has to copy the
# entire image data. So, on a large image, it becomes super slow.
# Therefore we cache the array. The array is marked read-only so we
# don't have to worry about it and the PIL image getting out of sync
# with modifications.
if self._array is None:
self._array = np.asarray(self._pil)
return self._array
[docs] def aspil(self):
"""Obtain the image data as :class:`PIL.Image.Image`.
Returns
-------
If the image was loaded as a PIL image, the underlying object will be
returned. Otherwise the data array will be converted into a PIL image,
which requires that the array have an RGB(A) format with a shape of
``(height, width, planes)``, where ``planes`` is 3 or 4, and a dtype of
``uint8``.
"""
if self._pil is not None:
return self._pil
return pil_image.fromarray(self._array)
@property
def mode(self):
return self._mode
@property
def dtype(self):
# TODO: can this be more efficient? Does it need to be?
return self.asarray().dtype
@property
def shape(self):
if self._array is not None:
return self._array.shape
return (self._pil.height, self._pil.width, len(self._pil.getbands()))
@property
def width(self):
return self.shape[1]
@property
def height(self):
return self.shape[0]
[docs] def fill_into_maskable_buffer(self, buffer, iy_idx, ix_idx, by_idx, bx_idx):
"""
Fill a maskable buffer with a rectangle of data from this image.
Parameters
----------
buffer : :class:`Image`
The destination buffer image, created with :meth:`ImageMode.make_maskable_buffer`.
iy_idx : slice or other indexer
The indexer into the Y axis of the source image (self).
ix_idx : slice or other indexer
The indexer into the X axis of the source image (self).
by_idx : slice or other indexer
The indexer into the Y axis of the destination *buffer*.
bx_idx : slice or other indexer
The indexer into the X axis of the destination *buffer*.
Notes
-----
This highly specialized function is used to tile images efficiently. No
bounds checking is performed. The rectangles defined by the indexers in
the source and destination are assumed to agree in size. The regions of
the buffer not filled by source data are masked, namely: either filled
with alpha=0 or with NaN, depending on the image mode.
"""
i = self.asarray()
b = buffer.asarray()
if self.mode == ImageMode.RGB:
b.fill(0)
b[by_idx,bx_idx,:3] = i[iy_idx,ix_idx]
b[by_idx,bx_idx,3] = 255
elif self.mode == ImageMode.RGBA:
b.fill(0)
b[by_idx,bx_idx] = i[iy_idx,ix_idx]
elif self.mode == ImageMode.F32:
b.fill(np.nan)
b[by_idx,bx_idx] = i[iy_idx,ix_idx]
else:
raise Exception('unhandled mode in fill_into_maskable_buffer')
[docs] def save_default(self, path_or_stream):
"""
Save this image to a filesystem path or stream
Parameters
----------
path_or_stream : path-like object or file-like object
The destination into which the data should be written. If file-like,
the stream should accept bytes.
"""
if self.mode in (ImageMode.RGB, ImageMode.RGBA):
self.aspil().save(path_or_stream, format='PNG')
elif self.mode == ImageMode.F32:
np.save(path_or_stream, self.asarray())
else:
raise Exception('unhandled mode in save_default')
[docs] def make_thumbnail_bitmap(self):
"""Create a thumbnail bitmap from the image.
Returns
-------
An RGB :class:`PIL.Image.Image` representing a thumbnail of the input
image. WWT thumbnails are 96 pixels wide and 45 pixels tall and should
be saved in JPEG format.
"""
if self.mode == ImageMode.F32:
raise Exception('cannot thumbnail-ify non-RGB Image')
THUMB_SHAPE = (96, 45)
THUMB_ASPECT = THUMB_SHAPE[0] / THUMB_SHAPE[1]
if self.width / self.height > THUMB_ASPECT:
# The image is wider than desired; we'll need to crop off the sides.
target_width = int(round(self.height * THUMB_ASPECT))
dx = (self.width - target_width) // 2
crop_box = (dx, 0, dx + target_width, self.height)
else:
# The image is taller than desired; crop off top and bottom.
target_height = int(round(self.width / THUMB_ASPECT))
dy = (self.height - target_height) // 2
crop_box = (0, dy, self.width, dy + target_height)
# Turns out that PIL the decompression-bomb checks happen here too.
# Disable in our usual, not thread-safe, way.
old_max = pil_image.MAX_IMAGE_PIXELS
try:
pil_image.MAX_IMAGE_PIXELS = None
thumb = self.aspil().crop(crop_box)
finally:
pil_image.MAX_IMAGE_PIXELS = old_max
thumb.thumbnail(THUMB_SHAPE)
# Depending on the source image, the mode might be RGBA, which can't
# be JPEG-ified.
thumb = thumb.convert('RGB')
return thumb
[docs] def clear(self):
"""
Fill the image with whatever "empty" value is most appropriate for its
mode.
Notes
-----
The image is assumed to be writable, which will not be the case for
images constructed from PIL. If the mode is RGB or RGBA, the buffer is
filled with zeros. If the mode is floating-point, the buffer is filled
with NaNs.
"""
if self._mode in (ImageMode.RGB, ImageMode.RGBA):
self.asarray().fill(0)
elif self._mode == ImageMode.F32:
self.asarray().fill(np.nan)
else:
raise Exception('unhandled mode in clear()')
