Spaces:
Runtime error
Runtime error
| # pylint: disable=relative-beyond-top-level | |
| from __future__ import annotations | |
| from typing import List, Optional, Union | |
| import numpy as np | |
| import navi | |
| from nodes.base_input import BaseInput, ErrorValue | |
| from ...impl.color.color import Color | |
| from ...utils.format import format_color_with_channels, format_image_with_channels | |
| from ...utils.utils import get_h_w_c | |
| class AudioInput(BaseInput): | |
| """Input a 1D Audio NumPy array""" | |
| def __init__(self, label: str = "Audio"): | |
| super().__init__("Audio", label) | |
| class ImageInput(BaseInput): | |
| """Input a 2D Image NumPy array""" | |
| def __init__( | |
| self, | |
| label: str = "Image", | |
| image_type: navi.ExpressionJson = "Image | Color", | |
| channels: Union[int, List[int], None] = None, | |
| allow_colors: bool = False, | |
| ): | |
| base_type = [navi.Image(channels=channels)] | |
| if allow_colors: | |
| base_type.append(navi.Color(channels=channels)) | |
| image_type = navi.intersect(image_type, base_type) | |
| super().__init__(image_type, label) | |
| self.channels: Optional[List[int]] = ( | |
| [channels] if isinstance(channels, int) else channels | |
| ) | |
| self.allow_colors: bool = allow_colors | |
| self.associated_type = np.ndarray | |
| if self.allow_colors: | |
| self.associated_type = Union[np.ndarray, Color] | |
| def enforce(self, value): | |
| if isinstance(value, Color): | |
| if not self.allow_colors: | |
| raise ValueError( | |
| f"The input {self.label} does not accept colors, but was connected with one." | |
| ) | |
| if self.channels is not None and value.channels not in self.channels: | |
| expected = format_color_with_channels(self.channels, plural=True) | |
| actual = format_color_with_channels([value.channels]) | |
| raise ValueError( | |
| f"The input {self.label} only supports {expected} but was given {actual}." | |
| ) | |
| return value | |
| assert isinstance(value, np.ndarray) | |
| _, _, c = get_h_w_c(value) | |
| if self.channels is not None and c not in self.channels: | |
| expected = format_image_with_channels(self.channels, plural=True) | |
| actual = format_image_with_channels([c]) | |
| raise ValueError( | |
| f"The input {self.label} only supports {expected} but was given {actual}." | |
| ) | |
| assert value.dtype == np.float32, "Expected the input image to be normalized." | |
| if c == 1 and value.ndim == 3: | |
| value = value[:, :, 0] | |
| return value | |
| def get_error_value(self, value) -> ErrorValue: | |
| def get_channels(channel: int) -> str: | |
| if channel == 1: | |
| return "Grayscale" | |
| if channel == 3: | |
| return "RGB" | |
| if channel == 4: | |
| return "RGBA" | |
| return f"{channel}-channel" | |
| if isinstance(value, Color): | |
| return { | |
| "type": "formatted", | |
| "formatString": f"{get_channels(value.channels)} Color", | |
| } | |
| elif isinstance(value, np.ndarray): | |
| h, w, c = get_h_w_c(value) | |
| return { | |
| "type": "formatted", | |
| "formatString": f"{get_channels(c)} Image {w}x{h}", | |
| } | |
| else: | |
| return super().get_error_value(value) | |
| class VideoInput(BaseInput): | |
| """Input a 3D Video NumPy array""" | |
| def __init__(self, label: str = "Video"): | |
| super().__init__("Video", label) | |