# -*- coding: utf-8 -*-

#

# Licensed under the terms of the BSD 3-Clause

# (see plotpy/LICENSE for details)

# pylint: disable=C0103

"""

Input/Output helper functions

-----------------------------

Overview

^^^^^^^^

The `io` module provides the following input/output helper functions:

* :py:func:`.io.imread`: load an image (.png, .tiff,

.dicom, etc.) and return its data as a NumPy array

* :py:func:`.io.imwrite`: save an array to an image file

* :py:func:`.io.load_items`: load plot items from HDF5

* :py:func:`.io.save_items`: save plot items to HDF5

Reference

^^^^^^^^^

.. autofunction:: imread

.. autofunction:: imwrite

.. autofunction:: load_items

.. autofunction:: save_items

"""

from __future__ import annotations

import logging

import os.path as osp

import re

import sys

from typing import TYPE_CHECKING, Any

import numpy as np

import PIL.Image

import PIL.TiffImagePlugin # py2exe

from plotpy.config import _

if TYPE_CHECKING:

import guidata.io

def scale_data_to_dtype(data: np.ndarray, dtype: np.dtype) -> np.ndarray:

"""Scale data to fit datatype dynamic range

Args:

data: data to scale

dtype: target datatype

Returns:

scaled data

.. warning::

This function modifies data in place

"""

info = np.iinfo(dtype)

dmin = data.min()

dmax = data.max()

data -= dmin

data = data * float(info.max - info.min) / (dmax - dmin)

data = data + float(info.min)

return np.array(data, dtype)

# ===============================================================================

# I/O File type definitions

# ===============================================================================

class FileType:

"""Filetype object:

* `name` : description of filetype,

* `read_func`, `write_func` : I/O callbacks,

* `extensions`: filename extensions (with a dot!) or filenames,

(list, tuple or space-separated string)

* `data_types`: supported data types"""

def __init__(

self,

name,

extensions,

read_func=None,

write_func=None,

data_types=None,

requires_template=False,

):

self.name = name

if isinstance(extensions, str):

extensions = extensions.split()

self.extensions = [osp.splitext(" " + ext)[1] for ext in extensions]

self.read_func = read_func

self.write_func = write_func

self.data_types = data_types

self.requires_template = requires_template

def matches(self, action, dtype, template):

"""Return True if file type matches passed data type and template

(or if dtype is None)"""

assert action in ("load", "save")

matches = dtype is None or self.data_types is None or dtype in self.data_types

if action == "save" and self.requires_template:

matches = matches and template is not None

return matches

@property

def wcards(self):

"""

:return:

"""

return "*" + (" *".join(self.extensions))

def filters(self, action, dtype, template):

"""

:param action:

:param dtype:

:param template:

:return:

"""

assert action in ("load", "save")

if self.matches(action, dtype, template):

return f"\n{self.name} ({self.wcards})"

else:

return ""

class ImageIOHandler:

"""I/O handler: regroup all FileType objects"""

def __init__(self):

self.filetypes = []

def allfilters(self, action, dtype, template):

"""

:param action:

:param dtype:

:param template:

:return:

"""

wcards = " ".join(

[

ftype.wcards

for ftype in self.filetypes

if ftype.matches(action, dtype, template)

]

)

return "{} ({})".format(_("All supported files"), wcards)

def get_filters(self, action, dtype=None, template=None):

"""Return file type filters for `action` (string: 'save' or 'load'),

`dtype` data type (None: all data types), and `template` (True if save

function requires a template (e.g. DICOM files), False otherwise)"""

filters = self.allfilters(action, dtype, template)

for ftype in self.filetypes:

filters += ftype.filters(action, dtype, template)

return filters

def add(

self,

name,

extensions,

read_func=None,

write_func=None,

import_func=None,

data_types=None,

requires_template=None,

):

"""

:param name:

:param extensions:

:param read_func:

:param write_func:

:param import_func:

:param data_types:

:param requires_template:

:return:

"""

if import_func is not None:

try:

import_func()

except ImportError:

return

assert read_func is not None or write_func is not None

ftype = FileType(

name,

extensions,

read_func=read_func,

write_func=write_func,

data_types=data_types,

requires_template=requires_template,

)

self.filetypes.append(ftype)

def _get_filetype(self, ext):

"""Return FileType object associated to file extension `ext`"""

for ftype in self.filetypes:

if ext.lower() in ftype.extensions:

return ftype

else:

raise RuntimeError(f"Unsupported file type: '{ext}'")

def get_readfunc(self, ext):

"""Return read function associated to file extension `ext`"""

ftype = self._get_filetype(ext)

if ftype.read_func is None:

raise RuntimeError(f"Unsupported file type (read): '{ext}'")

else:

return ftype.read_func

def get_writefunc(self, ext):

"""Return read function associated to file extension `ext`"""

ftype = self._get_filetype(ext)

if ftype.write_func is None:

raise RuntimeError(f"Unsupported file type (write): '{ext}'")

else:

return ftype.write_func

iohandler = ImageIOHandler()

# ==============================================================================

# tifffile-based Private I/O functions

# ==============================================================================

def _imread_tiff(filename):

"""Open a TIFF image and return a NumPy array"""

try:

import tifffile

return tifffile.imread(filename)

except ImportError:

return _imread_pil(filename)

def _imwrite_tiff(filename, arr):

"""Save a NumPy array to a TIFF file"""

try:

import tifffile

return tifffile.imwrite(filename, arr)

except ImportError:

return _imwrite_pil(filename, arr)

# ==============================================================================

# PIL-based Private I/O functions

# ==============================================================================

if sys.byteorder == "little":

_ENDIAN = "<"

else:

_ENDIAN = ">"

DTYPES = {

"1": ("|b1", None),

"L": ("|u1", None),

"I": ("%si4" % _ENDIAN, None),

"F": ("%sf4" % _ENDIAN, None),

"I;16": ("%su2" % _ENDIAN, None),

"I;16B": ("%su2" % _ENDIAN, None),

"I;16S": ("%si2" % _ENDIAN, None),

"P": ("|u1", None),

"RGB": ("|u1", 3),

"RGBX": ("|u1", 4),

"RGBA": ("|u1", 4),

"CMYK": ("|u1", 4),

"YCbCr": ("|u1", 4),

}

def _imread_pil(filename, to_grayscale=False, **kwargs):

"""Open image with PIL and return a NumPy array"""

PIL.TiffImagePlugin.OPEN_INFO[(PIL.TiffImagePlugin.II, 0, 1, 1, (16,), ())] = (

"I;16",

"I;16",

)

with PIL.Image.open(filename) as img:

base, ext = osp.splitext(filename)

ext = ext.lower()

if ext in [".tif", ".tiff"]:

# try to know if multiple pages

nb_pages = 0

while True:

try:

img.seek(nb_pages)

nb_pages += 1

except EOFError:

break

if nb_pages > 1:

for i in range(nb_pages):

img.seek(i)

filename = base

filename += "_{i:d}".format(i=i)

filename += ext

img.save(filename)

if nb_pages == 2:

# possibility to be TIFF file with thumbnail and full image

# --> try to load full image (second one)

filename = base + "_{i:d}".format(i=1) + ext

else:

# we don't know which one must be loaded --> load first image

filename = base + "_{i:d}".format(i=0) + ext

with PIL.Image.open(filename) as img:

if img.mode in ("CMYK", "YCbCr"):

# Converting to RGB

img = img.convert("RGB")

if to_grayscale and img.mode in ("RGB", "RGBA", "RGBX"):

# Converting to grayscale

img = img.convert("L")

elif "A" in img.mode or (img.mode == "P" and "transparency" in img.info):

img = img.convert("RGBA")

elif img.mode == "P":

img = img.convert("RGB")

try:

dtype, extra = DTYPES[img.mode]

except KeyError:

raise RuntimeError(f"{img.mode} mode is not supported")

shape = (img.size[1], img.size[0])

if extra is not None:

shape += (extra,)

try:

return np.array(img, dtype=np.dtype(dtype)).reshape(shape)

except SystemError:

return np.array(img.getdata(), dtype=np.dtype(dtype)).reshape(shape)

def _imwrite_pil(filename, arr):

"""Write `arr` NumPy array to `filename` using PIL"""

for mode, (dtype_str, extra) in list(DTYPES.items()):

if dtype_str == arr.dtype.str:

if extra is None: # mode for grayscale images

if len(arr.shape[2:]) > 0:

continue # not suitable for RGB(A) images

else:

break # this is it!

else: # mode for RGB(A) images

if len(arr.shape[2:]) == 0:

continue # not suitable for grayscale images

elif arr.shape[-1] == extra:

break # this is it!

else:

# F Chermette 2022

if arr.dtype.str == "%sf8" % _ENDIAN:

arr = np.array(arr, dtype="f4")

mode = "F"

else:

raise RuntimeError("Cannot determine PIL data type")

img = PIL.Image.fromarray(arr, mode)

img.save(filename)

# ==============================================================================

# DICOM Private I/O functions

# ==============================================================================

def _import_dcm():

"""DICOM Import function (checking for required libraries):

DICOM support requires library `pydicom`"""

logger = logging.getLogger("pydicom")

logger.setLevel(logging.CRITICAL)

# This import statement must stay here because the purpose of this function

# is to check if pydicom is installed:

# pylint: disable=import-outside-toplevel

# pylint: disable=import-error

from pydicom import dcmread # type:ignore # noqa: F401

logger.setLevel(logging.WARNING)

def _imread_dcm(filename, **kwargs):

"""Open DICOM image with pydicom and return a NumPy array"""

# pylint: disable=import-outside-toplevel

# pylint: disable=import-error

from pydicom import dcmread # type:ignore

dcm = dcmread(filename, force=True)

# **********************************************************************

# The following is necessary until pydicom numpy support is improved:

# (after that, a simple: 'arr = dcm.PixelArray' will work the same)

format_str = "%sint%s" % (("u", "")[dcm.PixelRepresentation], dcm.BitsAllocated)

try:

dtype = np.dtype(format_str)

except TypeError:

raise TypeError(

"Data type not understood by NumPy: "

"PixelRepresentation=%d, BitsAllocated=%d"

% (dcm.PixelRepresentation, dcm.BitsAllocated)

)

arr = np.frombuffer(dcm.PixelData, dtype)

try:

# pydicom 0.9.3:

dcm_is_little_endian = dcm.isLittleEndian

except AttributeError:

# pydicom 0.9.4:

dcm_is_little_endian = dcm.is_little_endian

if dcm_is_little_endian != (sys.byteorder == "little"):

arr.byteswap(True)

spp = getattr(dcm, "SamplesperPixel", 1)

if hasattr(dcm, "NumberOfFrames") and dcm.NumberOfFrames > 1:

if spp > 1:

arr = arr.reshape(spp, dcm.NumberofFrames, dcm.Rows, dcm.Columns)

else:

arr = arr.reshape(dcm.NumberOfFrames, dcm.Rows, dcm.Columns)

else:

if spp > 1:

if dcm.BitsAllocated == 8:

arr = arr.reshape(spp, dcm.Rows, dcm.Columns)

else:

raise NotImplementedError(

"This code only handles SamplesPerPixel > 1 if Bits Allocated = 8"

)

else:

arr = arr.reshape(dcm.Rows, dcm.Columns)

# **********************************************************************

return arr

def _imwrite_dcm(filename, arr, template=None):

"""Save a numpy array `arr` into a DICOM image file `filename`

based on DICOM structure `template`"""

# Note: due to IOHandler formalism, `template` has to be a keyword argument

assert template is not None, (

"The `template` keyword argument is required to save DICOM files\n"

"(that is the template DICOM structure object)"

)

info = np.iinfo(arr.dtype)

template.BitsAllocated = info.bits

template.BitsStored = info.bits

template.HighBit = info.bits - 1

template.PixelRepresentation = ("u", "i").index(info.kind)

data_vr = ("US", "SS")[template.PixelRepresentation]

template.Rows = arr.shape[0]

template.Columns = arr.shape[1]

template.SmallestImagePixelValue = int(arr.min())

template[0x00280106].VR = data_vr

template.LargestImagePixelValue = int(arr.max())

template[0x00280107].VR = data_vr

if not template.PhotometricInterpretation.startswith("MONOCHROME"):

template.PhotometricInterpretation = "MONOCHROME1"

template.PixelData = arr.tostring()

template[0x7FE00010].VR = "OB"

template.save_as(filename)

# ==============================================================================

# Text files Private I/O functions

# ==============================================================================

def _imread_txt(filename, **kwargs):

"""Open text file image and return a NumPy array"""

for delimiter in ("\t", ",", " ", ";"):

try:

return np.loadtxt(filename, delimiter=delimiter)

except ValueError:

continue

else:

raise ValueError(f"Could not load {filename!r}")

def _imwrite_txt(filename, arr):

"""Write `arr` NumPy array to text file `filename`"""

if arr.dtype in (np.int8, np.uint8, np.int16, np.uint16, np.int32, np.uint32):

fmt = "%d"

else:

fmt = "%.18e"

ext = osp.splitext(filename)[1]

if ext.lower() in (".txt", ".asc", ""):

np.savetxt(filename, arr, fmt=fmt)

elif ext.lower() == ".csv":

np.savetxt(filename, arr, fmt=fmt, delimiter=",")

# ==============================================================================

# Registering I/O functions

# ==============================================================================

iohandler.add(

_("PNG files"),

"*.png",

read_func=_imread_pil,

write_func=_imwrite_pil,

data_types=(np.uint8, np.uint16),

)

iohandler.add(

_("TIFF files"), "*.tif *.tiff", read_func=_imread_tiff, write_func=_imwrite_tiff

)

iohandler.add(

_("8-bit images"),

"*.jpg *.gif",

read_func=_imread_pil,

write_func=_imwrite_pil,

data_types=(np.uint8,),

)

iohandler.add(_("NumPy arrays"), "*.npy", read_func=np.load, write_func=np.save)

iohandler.add(

_("Text files"), "*.txt *.csv *.asc", read_func=_imread_txt, write_func=_imwrite_txt

)

iohandler.add(

_("DICOM files"),

"*.dcm",

read_func=_imread_dcm,

write_func=_imwrite_dcm,

import_func=_import_dcm,

data_types=(np.int8, np.uint8, np.int16, np.uint16),

requires_template=True,

)

# ==============================================================================

# Generic image read/write functions

# ==============================================================================

def imread(

fname: str, ext: str | None = None, to_grayscale: bool = False

) -> np.ndarray:

"""Read an image from a file as a NumPy array

Args:

fname: image filename

ext: image file extension (if None, extension is guessed from filename)

to_grayscale: convert RGB images to grayscale

"""

if not isinstance(fname, str):

fname = str(fname) # in case filename is a QString instance

if ext is None:

_base, ext = osp.splitext(fname)

arr = iohandler.get_readfunc(ext)(fname)

dtype = arr.dtype

if to_grayscale and arr.ndim == 3:

# Converting to grayscale

arr = arr[..., :4].mean(axis=2)

arr = arr.astype(dtype)

return arr

else:

return arr

def imwrite(

fname: str,

arr: np.ndarray,

ext: str | None = None,

dtype: np.dtype | None = None,

max_range: bool | None = None,

**kwargs,

) -> None:

"""Write a NumPy array to an image file

Args:

fname: image filename

arr: NumPy array

ext: image file extension (if None, extension is guessed from filename)

dtype: data type (if None, data type is guessed from array)

max_range: scale data to fit dtype dynamic range

kwargs: additional keyword arguments passed to the image writer

.. warning::

If `max_range` is True, array data is modified in place

"""

if not isinstance(fname, str):

fname = str(fname) # in case filename is a QString instance

if ext is None:

_base, ext = osp.splitext(fname)

if max_range:

arr = scale_data_to_dtype(arr, arr.dtype if dtype is None else dtype)

iohandler.get_writefunc(ext)(fname, arr, **kwargs)

# ==============================================================================

# plotpy plot items I/O

# ==============================================================================

SERIALIZABLE_ITEMS = []

ITEM_MODULES = {}

def register_serializable_items(modname, classnames):

"""Register serializable item from module name and class name"""

global SERIALIZABLE_ITEMS, ITEM_MODULES

SERIALIZABLE_ITEMS += classnames

ITEM_MODULES[modname] = ITEM_MODULES.setdefault(modname, []) + classnames

# Curves

register_serializable_items(

"plotpy.items",

[

"CurveItem",

"PolygonMapItem",

"ErrorBarCurveItem",

"RawImageItem",

"ImageItem",

"XYImageItem",

"RGBImageItem",

"TrImageItem",

"MaskedImageItem",

"MaskedXYImageItem",

"Marker",

"XRangeSelection",

"YRangeSelection",

"PolygonShape",

"PointShape",

"SegmentShape",

"RectangleShape",

"ObliqueRectangleShape",

"EllipseShape",

"Axes",

"AnnotatedPoint",

"AnnotatedSegment",

"AnnotatedXRange",

"AnnotatedYRange",

"AnnotatedRectangle",

"AnnotatedObliqueRectangle",

"AnnotatedEllipse",

"AnnotatedCircle",

"AnnotatedPolygon",

"LabelItem",

"LegendBoxItem",

"SelectedLegendBoxItem",

],

)

def item_class_from_name(name: str) -> type[Any] | None:

"""Return plot item class from class name

Args:

name: plot item class name

Returns:

plot item class

Raises:

AssertionError: if class name is unknown (item is not serializable)

"""

global SERIALIZABLE_ITEMS, ITEM_MODULES

assert name in SERIALIZABLE_ITEMS, "Unknown class %r" % name

for modname, names in list(ITEM_MODULES.items()):

if name in names:

return getattr(__import__(modname, fromlist=[name]), name)

def item_name_from_object(obj: Any) -> str | None:

"""Return plot item class name from instance

Args:

obj: plot item instance

Returns:

plot item class name

"""

return obj.__class__.__name__

def save_item(

writer: guidata.io.HDF5Writer | guidata.io.INIWriter | guidata.io.JSONWriter,

group_name,

item: Any,

) -> None:

"""Save plot item to HDF5, INI or JSON file

Args:

writer: HDF5, INI or JSON writer

group_name: group name

item: serializable plot item

"""

with writer.group(group_name):

if item is None:

writer.write_none()

else:

item.serialize(writer)

with writer.group("item_class_name"):

writer.write_str(item_name_from_object(item))

def load_item(

reader: guidata.io.HDF5Reader | guidata.io.INIReader | guidata.io.JSONReader,

group_name,

) -> Any | None:

"""Load plot item from HDF5, INI or JSON file

Args:

reader: HDF5, INI or JSON reader

group_name: group name

Returns:

Plot item instance

"""

with reader.group(group_name):

with reader.group("item_class_name"):

try:

klass_name = reader.read_str()

except ValueError:

# None was saved instead of a real item

return

klass = item_class_from_name(klass_name)

item = klass()

item.deserialize(reader)

return item

def save_items(

writer: guidata.io.HDF5Writer | guidata.io.INIWriter | guidata.io.JSONWriter,

items: list[Any],

) -> None:

"""Save items to HDF5, INI or JSON file

Args:

writer: HDF5, INI or JSON writer

items: list of serializable plot items

"""

counts = {}

names = []

def _get_name(item):

basename = item_name_from_object(item)

count = counts[basename] = counts.setdefault(basename, 0) + 1

name = "%s_%03d" % (basename, count)

names.append(name.encode("utf-8"))

return name

for item in items:

with writer.group(_get_name(item)):

item.serialize(writer)

writer.write(item.isVisible(), group_name="visible")

with writer.group("plot_items"):

writer.write_sequence(names)

def load_items(

reader: guidata.io.HDF5Reader | guidata.io.INIReader | guidata.io.JSONReader,

) -> list[Any]:

"""Load items from HDF5, INI or JSON file

Args:

reader: HDF5, INI or JSON reader

Returns:

list of plot item instances

"""

with reader.group("plot_items"):

names = reader.read_sequence()

items = []

for name in names:

try:

name_str = name.decode()

except AttributeError:

name_str = name

klass_name = re.match(r"([A-Z]+[A-Za-z0-9\_]*)\_([0-9]*)", name_str).groups()[0]

klass = item_class_from_name(klass_name)

item = klass()

with reader.group(name):

item.deserialize(reader)

item.setVisible(reader.read("visible", default=True))

items.append(item)

return items

if __name__ == "__main__":

# Test if items can all be constructed from their Python module

for name in SERIALIZABLE_ITEMS:

print(name, "-->", item_class_from_name(name))