#!/usr/bin/env python3 """ Common multi-backend Keras utilities """ from __future__ import annotations import typing as T import numpy as np from keras import ops, Variable from lib.utils import get_module_objects if T.TYPE_CHECKING: from keras import KerasTensor # TODO these can probably be switched to pure pytorch def frobenius_norm(matrix: KerasTensor, axis: int = -1, keep_dims: bool = True, epsilon: float = 1e-15) -> KerasTensor: """ Frobenius normalization for Keras Tensor Parameters ---------- matrix: :class:`keras.KerasTensor` The matrix to normalize axis: int, optional The axis to normalize. Default: `-1` keep_dims: bool, Optional Whether to retain the original matrix shape or not. Default:``True`` epsilon: flot, optional Epsilon to apply to the normalization to preven NaN errors on zero values Returns ------- :class:`keras.KerasTensor` The normalized output """ return ops.sqrt(ops.sum(ops.power(matrix, 2), axis=axis, keepdims=keep_dims) + epsilon) def replicate_pad(image: KerasTensor, padding: int) -> KerasTensor: """ Apply replication padding to an input batch of images. Expects 4D tensor in BHWC format. Notes ----- At the time of writing Keras does not have a native replication padding method. The implementation here is probably not the most efficient, but it is a pure keras method which should work ok. Parameters ---------- image: :class:`keras.KerasTensor` Image tensor to pad pad: int The amount of padding to apply to each side of the input image Returns ------- :class:`keras.KerasTensor` The input image with replication padding applied """ top_pad = ops.tile(image[:, :1, ...], (1, padding, 1, 1)) bottom_pad = ops.tile(image[:, -1:, ...], (1, padding, 1, 1)) pad_top_bottom = ops.concatenate([top_pad, image, bottom_pad], axis=1) left_pad = ops.tile(pad_top_bottom[..., :1, :], (1, 1, padding, 1)) right_pad = ops.tile(pad_top_bottom[..., -1:, :], (1, 1, padding, 1)) padded = ops.concatenate([left_pad, pad_top_bottom, right_pad], axis=2) return padded class ColorSpaceConvert(): """ Transforms inputs between different color spaces on the GPU Notes ----- The following color space transformations are implemented: - rgb to lab - rgb to xyz - srgb to _rgb - srgb to ycxcz - xyz to ycxcz - xyz to lab - xyz to rgb - ycxcz to rgb - ycxcz to xyz Parameters ---------- from_space: str One of `"srgb"`, `"rgb"`, `"xyz"` to_space: str One of `"lab"`, `"rgb"`, `"ycxcz"`, `"xyz"` Raises ------ ValueError If the requested color space conversion is not defined """ def __init__(self, from_space: str, to_space: str) -> None: functions = {"rgb_lab": self._rgb_to_lab, "rgb_xyz": self._rgb_to_xyz, "srgb_rgb": self._srgb_to_rgb, "srgb_ycxcz": self._srgb_to_ycxcz, "xyz_ycxcz": self._xyz_to_ycxcz, "xyz_lab": self._xyz_to_lab, "xyz_rgb": self._xyz_to_rgb, "ycxcz_rgb": self._ycxcz_to_rgb, "ycxcz_xyz": self._ycxcz_to_xyz} func_name = f"{from_space.lower()}_{to_space.lower()}" if func_name not in functions: raise ValueError(f"The color transform {from_space} to {to_space} is not defined.") self._func = functions[func_name] self._ref_illuminant = Variable(np.array([[[0.950428545, 1.000000000, 1.088900371]]]), dtype="float32", trainable=False) self._inv_ref_illuminant = 1. / self._ref_illuminant self._rgb_xyz_map = self._get_rgb_xyz_map() self._xyz_multipliers = Variable([116, 500, 200], dtype="float32", trainable=False) @classmethod def _get_rgb_xyz_map(cls) -> tuple[KerasTensor, KerasTensor]: """ Obtain the mapping and inverse mapping for rgb to xyz color space conversion. Returns ------- tuple The mapping and inverse Tensors for rgb to xyz color space conversion """ mapping = np.array([[10135552 / 24577794, 8788810 / 24577794, 4435075 / 24577794], [2613072 / 12288897, 8788810 / 12288897, 887015 / 12288897], [1425312 / 73733382, 8788810 / 73733382, 70074185 / 73733382]]) inverse = np.linalg.inv(mapping) return (Variable(mapping, dtype="float32", trainable=False), Variable(inverse, dtype="float32", trainable=False)) def __call__(self, image: KerasTensor) -> KerasTensor: """ Call the colorspace conversion function. Parameters ---------- image: :class:`keras.KerasTensor` The image tensor in the colorspace defined by :attr:`from_space` Returns ------- :class:`keras.KerasTensor` The image tensor in the colorspace defined by :attr:`to_space` """ return self._func(image) def _rgb_to_lab(self, image: KerasTensor) -> KerasTensor: """ RGB to LAB conversion. Parameters ---------- image: :class:`keras.KerasTensor` The image tensor in RGB format Returns ------- :class:`keras.KerasTensor` The image tensor in LAB format """ converted = self._rgb_to_xyz(image) return self._xyz_to_lab(converted) def _rgb_xyz_rgb(self, image: KerasTensor, mapping: KerasTensor) -> KerasTensor: """ RGB to XYZ or XYZ to RGB conversion. Notes ----- The conversion in both directions is the same, but the mappping matrix for XYZ to RGB is the inverse of RGB to XYZ. References ---------- https://www.image-engineering.de/library/technotes/958-how-to-convert-between-srgb-and-ciexyz Parameters ---------- mapping: :class:`keras.KerasTensor` The mapping matrix to perform either the XYZ to RGB or RGB to XYZ color space conversion image: :class:`keras.KerasTensor` The image tensor in RGB format Returns ------- :class:`keras.KerasTensor` The image tensor in XYZ format """ dim = image.shape image = ops.transpose(image, (0, 3, 1, 2)) image = ops.reshape(image, (dim[0], dim[3], dim[1] * dim[2])) converted = ops.transpose(ops.dot(mapping, image), (0, 2, 1)) return ops.reshape(converted, dim) def _rgb_to_xyz(self, image: KerasTensor) -> KerasTensor: """ RGB to XYZ conversion. Parameters ---------- image: :class:`keras.KerasTensor` The image tensor in RGB format Returns ------- :class:`keras.KerasTensor` The image tensor in XYZ format """ return self._rgb_xyz_rgb(image, self._rgb_xyz_map[0]) @classmethod def _srgb_to_rgb(cls, image: KerasTensor) -> KerasTensor: """ SRGB to RGB conversion. Notes ----- RGB Image is clipped to a small epsilon to stabalize training Parameters ---------- image: :class:`keras.KerasTensor` The image tensor in SRGB format Returns ------- :class:`keras.KerasTensor` The image tensor in RGB format """ limit = np.float32(0.04045) return ops.where(image > limit, ops.power((ops.clip(image, limit, np.inf) + 0.055) / 1.055, 2.4), image / 12.92) def _srgb_to_ycxcz(self, image: KerasTensor) -> KerasTensor: """ SRGB to YcXcZ conversion. Parameters ---------- image: :class:`keras.KerasTensor` The image tensor in SRGB format Returns ------- :class:`keras.KerasTensor` The image tensor in YcXcZ format """ converted = self._srgb_to_rgb(image) converted = self._rgb_to_xyz(converted) return self._xyz_to_ycxcz(converted) def _xyz_to_lab(self, image: KerasTensor) -> KerasTensor: """ XYZ to LAB conversion. Parameters ---------- image: :class:`keras.KerasTensor` The image tensor in XYZ format Returns ------- :class:`keras.KerasTensor` The image tensor in LAB format """ image = image * self._inv_ref_illuminant delta = 6 / 29 delta_cube = delta ** 3 factor = 1 / (3 * (delta ** 2)) clamped_term = ops.power(ops.clip(image, delta_cube, np.inf), 1.0 / 3.0) div = factor * image + (4 / 29) image = ops.where(image > delta_cube, clamped_term, div) return ops.concatenate([self._xyz_multipliers[0] * image[..., 1:2] - 16., self._xyz_multipliers[1:] * (image[..., :2] - image[..., 1:3])], axis=-1) def _xyz_to_rgb(self, image: KerasTensor) -> KerasTensor: """ XYZ to YcXcZ conversion. Parameters ---------- image: :class:`keras.KerasTensor` The image tensor in XYZ format Returns ------- :class:`keras.KerasTensor` The image tensor in RGB format """ return self._rgb_xyz_rgb(image, self._rgb_xyz_map[1]) def _xyz_to_ycxcz(self, image: KerasTensor) -> KerasTensor: """ XYZ to YcXcZ conversion. Parameters ---------- image: :class:`keras.KerasTensor` The image tensor in XYZ format Returns ------- :class:`keras.KerasTensor` The image tensor in YcXcZ format """ image = image * self._inv_ref_illuminant return ops.concatenate([self._xyz_multipliers[0] * image[..., 1:2] - 16., self._xyz_multipliers[1:] * (image[..., :2] - image[..., 1:3])], axis=-1) def _ycxcz_to_rgb(self, image: KerasTensor) -> KerasTensor: """ YcXcZ to RGB conversion. Parameters ---------- image: :class:`keras.KerasTensor` The image tensor in YcXcZ format Returns ------- :class:`keras.KerasTensor` The image tensor in RGB format """ converted = self._ycxcz_to_xyz(image) return self._xyz_to_rgb(converted) def _ycxcz_to_xyz(self, image: KerasTensor) -> KerasTensor: """ YcXcZ to XYZ conversion. Parameters ---------- image: :class:`keras.KerasTensor` The image tensor in YcXcZ format Returns ------- :class:`keras.KerasTensor` The image tensor in XYZ format """ ch_y = (image[..., 0:1] + 16.) / self._xyz_multipliers[0] return ops.concatenate([ch_y + (image[..., 1:2] / self._xyz_multipliers[1]), ch_y, ch_y - (image[..., 2:3] / self._xyz_multipliers[2])], axis=-1) * self._ref_illuminant __all__ = get_module_objects(__name__)