Spaces:
Runtime error
Runtime error
File size: 2,142 Bytes
c19ca42 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 |
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)
|