Spaces:
Runtime error
Runtime error
| from enum import Enum | |
| from typing import List | |
| import numpy as np | |
| from ..color.convert import convert | |
| from ..color.convert_data import LAB, RGB | |
| from ..image_op import ImageOp | |
| class SplitMode(Enum): | |
| RGB = 1 | |
| LAB = 2 | |
| def split(self, img: np.ndarray) -> List[np.ndarray]: | |
| if img.ndim == 2: | |
| return [img] | |
| assert img.ndim == 3 | |
| c = img.shape[2] | |
| if c == 1: | |
| return [img[:, :, 0]] | |
| if self == SplitMode.RGB: | |
| return [img[:, :, channel] for channel in range(c)] | |
| elif self == SplitMode.LAB: | |
| if c < 3: | |
| return [img[:, :, channel] for channel in range(c)] | |
| lab = convert(img[:, :, 0:3], RGB, LAB) | |
| remaining_channels = [img[:, :, channel] for channel in range(3, c)] | |
| return [ | |
| lab[:, :, 0], | |
| lab[:, :, 1], | |
| lab[:, :, 2], | |
| *remaining_channels, | |
| ] | |
| else: | |
| raise AssertionError | |
| def combine(self, channels: List[np.ndarray]) -> np.ndarray: | |
| l = len(channels) | |
| assert l > 0 | |
| if l == 1: | |
| return channels[0] | |
| if self == SplitMode.RGB: | |
| return np.dstack(channels) | |
| elif self == SplitMode.LAB: | |
| if l < 3: | |
| return np.dstack(channels) | |
| rgb = convert(np.dstack(channels[0:3]), LAB, RGB) | |
| if l == 3: | |
| return rgb | |
| return np.dstack([rgb[:, :, 0], rgb[:, :, 1], rgb[:, :, 2], *channels[3:]]) | |
| else: | |
| raise AssertionError | |
| def grayscale_split( | |
| img: np.ndarray, process: ImageOp, mode: SplitMode = SplitMode.RGB | |
| ) -> np.ndarray: | |
| """ | |
| This function guarantees that the given image operation method will be called with 2D single-channel images. | |
| The images passed into the operation are guaranteed to have the same size as the given image. | |
| """ | |
| input_channels = mode.split(img) | |
| output_channels: List[np.ndarray] = [] | |
| for channel in input_channels: | |
| output_channels.append(process(channel)) | |
| return mode.combine(output_channels) | |