[Core] Add ResolveImage node
Copy-paste current implementation of core LoadImage, but use a string input providing an entity reference instead of a file upload, and retrieve input file path from the OpenAssetIO manager. Due to the mix-and-match by ComfyUI of class vs. instance methods, plus no way to inject context, use a singleton class to represent the OpenAssetIO host application. Signed-off-by: David Feltell <david.feltell@foundry.com>
This commit is contained in:
David Feltell
@@ -29,3 +29,6 @@ indent_style = tab
|
||||
|
||||
[{*.yml,*.yaml}]
|
||||
indent_size = 2
|
||||
|
||||
[*.json]
|
||||
indent_size = 2
|
||||
+15
-3
@@ -2,7 +2,16 @@
|
||||
# Copyright (c) 2025 The Foundry Visionmongers Ltd
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
"""Top-level package for openassetio-comfyui."""
|
||||
"""
|
||||
Top-level package for openassetio-comfyui.
|
||||
|
||||
This __init__.py will not be packaged, but is useful if the repository
|
||||
is checked out under ComfyUI's custom_nodes directory, as it allows
|
||||
ComfyUI to discover the nodes contained within.
|
||||
"""
|
||||
|
||||
import sys
|
||||
import pathlib
|
||||
|
||||
__all__ = [
|
||||
"NODE_CLASS_MAPPINGS",
|
||||
@@ -13,5 +22,8 @@ __author__ = """Contributors to the OpenAssetIO project"""
|
||||
__email__ = "openassetio-discussion@lists.aswf.io"
|
||||
__version__ = "1.0.0"
|
||||
|
||||
from .src.openassetio_comfyui.nodes import NODE_CLASS_MAPPINGS
|
||||
from .src.openassetio_comfyui.nodes import NODE_DISPLAY_NAME_MAPPINGS
|
||||
# Ensure src/ is on the path so we can import from there.
|
||||
sys.path.append(str(pathlib.Path(__file__).parent / "src"))
|
||||
|
||||
from openassetio_comfyui.nodes import NODE_CLASS_MAPPINGS
|
||||
from openassetio_comfyui.nodes import NODE_DISPLAY_NAME_MAPPINGS
|
||||
|
||||
+32
-4
@@ -17,8 +17,10 @@ readme = "README.md"
|
||||
license = { text = "Apache Software License 2.0" }
|
||||
requires-python = ">=3.10,<3.12"
|
||||
classifiers = []
|
||||
dependencies = []
|
||||
|
||||
dependencies = [
|
||||
"openassetio>=1.0.0",
|
||||
"openassetio-mediacreation>=1.0.0a12"
|
||||
]
|
||||
|
||||
[project.optional-dependencies]
|
||||
dev = [
|
||||
@@ -28,6 +30,7 @@ dev = [
|
||||
"pre-commit", # runs linting on commit
|
||||
"pytest", # testing
|
||||
"ruff", # linting
|
||||
"openassetio-manager-bal" # mock manager for testing
|
||||
]
|
||||
|
||||
[project.urls]
|
||||
@@ -77,12 +80,37 @@ target-version = "py311"
|
||||
# Add rules to ban exec/eval
|
||||
[tool.ruff.lint]
|
||||
select = [
|
||||
# See all rules here: https://docs.astral.sh/ruff/rules
|
||||
"S102", # exec-builtin
|
||||
"S307", # eval-used
|
||||
"W293",
|
||||
"F", # The "F" series in Ruff stands for "Pyflakes" rules, which catch various Python syntax errors and undefined names.
|
||||
# See all rules here: https://docs.astral.sh/ruff/rules/#pyflakes-f
|
||||
# The "F" series in Ruff stands for "Pyflakes" rules, which catch
|
||||
# various Python syntax errors and undefined names.
|
||||
"F",
|
||||
# pydocstyle
|
||||
"D",
|
||||
# pycodestyle
|
||||
"E", "W",
|
||||
# pylint
|
||||
"PL"
|
||||
]
|
||||
|
||||
ignore = [
|
||||
# "One-line docstring should fit on one line" - i.e. expects
|
||||
# """summary""", not """\nsummary\n""", which is inconsistent with
|
||||
# other OpenAssetIO projects.
|
||||
"D200",
|
||||
# "1 blank line required between summary line and description" -
|
||||
# triggers if the summary spans more than one line.
|
||||
"D205"
|
||||
]
|
||||
|
||||
[tool.ruff.lint.flake8-quotes]
|
||||
inline-quotes = "double"
|
||||
|
||||
[tool.ruff.lint.pydocstyle]
|
||||
convention = "pep257"
|
||||
|
||||
[tool.ruff.lint.pycodestyle]
|
||||
max-doc-length = 72
|
||||
max-line-length = 99
|
||||
|
||||
+235
-101
@@ -1,131 +1,265 @@
|
||||
# openassetio-comfyui
|
||||
# Copyright (c) 2025 The Foundry Visionmongers Ltd
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
"""
|
||||
OpenAssetIO nodes for ComfyUI.
|
||||
|
||||
from inspect import cleandoc
|
||||
Also contains a singleton class representing the OpenAssetIO host
|
||||
application.
|
||||
"""
|
||||
|
||||
import hashlib
|
||||
import logging
|
||||
|
||||
import torch
|
||||
import numpy as np
|
||||
|
||||
from PIL import Image, ImageOps, ImageSequence
|
||||
|
||||
from openassetio import EntityReference
|
||||
from openassetio.utils import FileUrlPathConverter
|
||||
from openassetio.log import LoggerInterface
|
||||
from openassetio.access import ResolveAccess
|
||||
from openassetio.hostApi import ManagerFactory, HostInterface
|
||||
from openassetio.pluginSystem import (
|
||||
HybridPluginSystemManagerImplementationFactory,
|
||||
PythonPluginSystemManagerImplementationFactory,
|
||||
CppPluginSystemManagerImplementationFactory,
|
||||
)
|
||||
from openassetio_mediacreation.traits.content import LocatableContentTrait
|
||||
|
||||
import node_helpers
|
||||
|
||||
|
||||
class Example:
|
||||
class _OpenAssetIOHost:
|
||||
"""
|
||||
A example node
|
||||
Singleton class representing the OpenAssetIO host.
|
||||
|
||||
Class methods
|
||||
-------------
|
||||
INPUT_TYPES (dict):
|
||||
Tell the main program input parameters of nodes.
|
||||
IS_CHANGED:
|
||||
optional method to control when the node is re executed.
|
||||
Instantiates and exposes a manager and context, and provides common
|
||||
utility methods.
|
||||
|
||||
Attributes
|
||||
----------
|
||||
RETURN_TYPES (`tuple`):
|
||||
The type of each element in the output tulple.
|
||||
RETURN_NAMES (`tuple`):
|
||||
Optional: The name of each output in the output tulple.
|
||||
FUNCTION (`str`):
|
||||
The name of the entry-point method. For example, if `FUNCTION = "execute"` then it will run Example().execute()
|
||||
OUTPUT_NODE ([`bool`]):
|
||||
If this node is an output node that outputs a result/image from the graph. The SaveImage node is an example.
|
||||
The backend iterates on these output nodes and tries to execute all their parents if their parent graph is properly connected.
|
||||
Assumed to be False if not present.
|
||||
CATEGORY (`str`):
|
||||
The category the node should appear in the UI.
|
||||
execute(s) -> tuple || None:
|
||||
The entry point method. The name of this method must be the same as the value of property `FUNCTION`.
|
||||
For example, if `FUNCTION = "execute"` then this method's name must be `execute`, if `FUNCTION = "foo"` then it must be `foo`.
|
||||
The OpenAssetIO manager must be configured via a config file,
|
||||
located via the OPENASSETIO_DEFAULT_CONFIG environment variable. See
|
||||
the OpenAssetIO documentation for more details.
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
pass
|
||||
_instance = None
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
def instance(cls) -> "_OpenAssetIOHost":
|
||||
"""
|
||||
Return a dictionary which contains config for all input fields.
|
||||
Some types (string): "MODEL", "VAE", "CLIP", "CONDITIONING", "LATENT", "IMAGE", "INT", "STRING", "FLOAT".
|
||||
Input types "INT", "STRING" or "FLOAT" are special values for fields on the node.
|
||||
The type can be a list for selection.
|
||||
Get or create the singleton instance of the OpenAssetIO host.
|
||||
"""
|
||||
if cls._instance is None:
|
||||
cls._instance = cls()
|
||||
return cls._instance
|
||||
|
||||
Returns: `dict`:
|
||||
- Key input_fields_group (`string`): Can be either required, hidden or optional. A node class must have property `required`
|
||||
- Value input_fields (`dict`): Contains input fields config:
|
||||
* Key field_name (`string`): Name of a entry-point method's argument
|
||||
* Value field_config (`tuple`):
|
||||
+ First value is a string indicate the type of field or a list for selection.
|
||||
+ Secound value is a config for type "INT", "STRING" or "FLOAT".
|
||||
class _HostInterface(HostInterface):
|
||||
def displayName(self) -> str:
|
||||
return "ComfyUI"
|
||||
|
||||
def identifier(self) -> str:
|
||||
return "org.foundry.comfyui"
|
||||
|
||||
def info(self) -> dict:
|
||||
return super().info()
|
||||
|
||||
class _LoggerInterface(LoggerInterface):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.__logger = logging.getLogger("openassetio-comfyui")
|
||||
|
||||
def log(self, severity, message) -> None:
|
||||
if severity == LoggerInterface.Severity.kDebug:
|
||||
self.__logger.debug(message)
|
||||
elif severity == LoggerInterface.Severity.kInfo:
|
||||
self.__logger.info(message)
|
||||
elif severity == LoggerInterface.Severity.kWarning:
|
||||
self.__logger.warning(message)
|
||||
elif severity == LoggerInterface.Severity.kError:
|
||||
self.__logger.error(message)
|
||||
elif severity == LoggerInterface.Severity.kCritical:
|
||||
self.__logger.critical(message)
|
||||
else:
|
||||
self.__logger.log(logging.NOTSET, message)
|
||||
|
||||
def __init__(self):
|
||||
"""
|
||||
Load and initialise an OpenAssetIO manager plugin.
|
||||
"""
|
||||
self.__logger = self._LoggerInterface()
|
||||
# Initialise plugin system, then find and load a manager.
|
||||
self.manager = ManagerFactory.defaultManagerForInterface(
|
||||
self._HostInterface(),
|
||||
HybridPluginSystemManagerImplementationFactory(
|
||||
# Prefer C++ over Python plugins/methods.
|
||||
[
|
||||
CppPluginSystemManagerImplementationFactory(self.__logger),
|
||||
PythonPluginSystemManagerImplementationFactory(self.__logger),
|
||||
],
|
||||
self.__logger,
|
||||
),
|
||||
self.__logger,
|
||||
)
|
||||
if self.manager is None:
|
||||
raise RuntimeError(
|
||||
"Could not create an OpenAssetIO manager instance. Ensure that your OpenAssetIO"
|
||||
" configuration is correct and that the environment variable"
|
||||
" OPENASSETIO_DEFAULT_CONFIG is set to a valid configuration file."
|
||||
)
|
||||
self.__context = self.manager.createContext()
|
||||
self.__file_url_path_converter = FileUrlPathConverter()
|
||||
|
||||
def resolve_to_path(self, entity_reference: str) -> str:
|
||||
"""
|
||||
Resolve an OpenAssetIO entity reference to a local file path.
|
||||
|
||||
OpenAssetIO will raise exceptions if the entity reference is
|
||||
invalid, if there is a problem resolving the entity, or if the
|
||||
given entity location is not a `file://` URL.
|
||||
|
||||
@throw RuntimeError if the entity reference resolves
|
||||
successfully, but the entity has no location.
|
||||
"""
|
||||
entity_reference = self.manager.createEntityReference(entity_reference)
|
||||
|
||||
traits_data = self.manager.resolve(
|
||||
entity_reference, {LocatableContentTrait.kId}, ResolveAccess.kRead, self.__context
|
||||
)
|
||||
locatable_content_trait = LocatableContentTrait(traits_data)
|
||||
url = locatable_content_trait.getLocation()
|
||||
|
||||
if url is None:
|
||||
raise RuntimeError(
|
||||
f"Failed to resolve entity reference '{entity_reference}' to a location"
|
||||
)
|
||||
|
||||
return self.__file_url_path_converter.pathFromUrl(url)
|
||||
|
||||
|
||||
class ResolveImage:
|
||||
"""
|
||||
Node to resolve an image from an OpenAssetIO entity reference.
|
||||
|
||||
The non-OpenAssetIO logic is largely duplicated from the built-in
|
||||
ComfyUI LoadImage node (at least as of 4f5812b9).
|
||||
"""
|
||||
|
||||
# Tooltip to display when hovering over the node.
|
||||
DESCRIPTION = "Resolve images from an asset manager."
|
||||
# Menu category.
|
||||
CATEGORY = "image"
|
||||
# Function to call when node is executed.
|
||||
FUNCTION = "resolve_image"
|
||||
# Node outputs.
|
||||
RETURN_TYPES = ("IMAGE", "MASK")
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls) -> dict:
|
||||
"""
|
||||
Input sockets and widgets.
|
||||
"""
|
||||
return {
|
||||
"required": {
|
||||
"image": ("Image", {"tooltip": "This is an image"}),
|
||||
"int_field": (
|
||||
"INT",
|
||||
{
|
||||
"default": 0,
|
||||
"min": 0, # Minimum value
|
||||
"max": 4096, # Maximum value
|
||||
"step": 64, # Slider's step
|
||||
"display": "number", # Cosmetic only: display as "number" or "slider"
|
||||
},
|
||||
),
|
||||
"float_field": (
|
||||
"FLOAT",
|
||||
{
|
||||
"default": 1.0,
|
||||
"min": 0.0,
|
||||
"max": 10.0,
|
||||
"step": 0.01,
|
||||
"round": 0.001, # The value represeting the precision to round to, will be set to the step value by default. Can be set to False to disable rounding.
|
||||
"display": "number",
|
||||
},
|
||||
),
|
||||
"print_to_screen": (["enable", "disable"],),
|
||||
"string_field": (
|
||||
"entity_reference": (
|
||||
"STRING",
|
||||
{
|
||||
"multiline": False, # True if you want the field to look like the one on the ClipTextEncode node
|
||||
"default": "Hello World!",
|
||||
"multiline": False,
|
||||
"default": "",
|
||||
"placeholder": "Entity reference...",
|
||||
"tooltip": "The entity to resolve",
|
||||
},
|
||||
),
|
||||
},
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
# RETURN_NAMES = ("image_output_name",)
|
||||
DESCRIPTION = cleandoc(__doc__)
|
||||
FUNCTION = "test"
|
||||
@classmethod
|
||||
def IS_CHANGED(cls, entity_reference: str) -> str:
|
||||
"""
|
||||
Resolve the entity reference to a file path, then return a hash
|
||||
of the file contents.
|
||||
|
||||
# OUTPUT_NODE = False
|
||||
# OUTPUT_TOOLTIPS = ("",) # Tooltips for the output node
|
||||
This allows ComfyUI to determine if the node needs to be
|
||||
re-executed.
|
||||
|
||||
CATEGORY = "Example"
|
||||
Non-OpenAssetIO logic largely duplicated from the built-in
|
||||
ComfyUI LoadImage node (at least as of 4f5812b9).
|
||||
"""
|
||||
image_path = _OpenAssetIOHost.instance().resolve_to_path(entity_reference)
|
||||
m = hashlib.sha256()
|
||||
with open(image_path, "rb") as f:
|
||||
m.update(f.read())
|
||||
return m.digest().hex()
|
||||
|
||||
def test(self, image, string_field, int_field, float_field, print_to_screen):
|
||||
if print_to_screen == "enable":
|
||||
print(f"""Your input contains:
|
||||
string_field aka input text: {string_field}
|
||||
int_field: {int_field}
|
||||
float_field: {float_field}
|
||||
""")
|
||||
# do some processing on the image, in this example I just invert it
|
||||
image = 1.0 - image
|
||||
return (image,)
|
||||
@classmethod
|
||||
def VALIDATE_INPUTS(cls, entity_reference: str) -> bool:
|
||||
"""
|
||||
Validate that the input entity reference is valid syntax for the
|
||||
current manager.
|
||||
"""
|
||||
return _OpenAssetIOHost.instance().manager.isEntityReferenceString(entity_reference)
|
||||
|
||||
"""
|
||||
The node will always be re executed if any of the inputs change but
|
||||
this method can be used to force the node to execute again even when the inputs don't change.
|
||||
You can make this node return a number or a string. This value will be compared to the one returned the last time the node was
|
||||
executed, if it is different the node will be executed again.
|
||||
This method is used in the core repo for the LoadImage node where they return the image hash as a string, if the image hash
|
||||
changes between executions the LoadImage node is executed again.
|
||||
"""
|
||||
# @classmethod
|
||||
# def IS_CHANGED(s, image, string_field, int_field, float_field, print_to_screen):
|
||||
# return ""
|
||||
def resolve_image(self, entity_reference: str) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
"""
|
||||
Resolve an OpenAssetIO entity reference to a file path, then
|
||||
load the image(s) at that path and return as tensors.
|
||||
|
||||
Non-OpenAssetIO logic largely duplicated from the built-in
|
||||
ComfyUI LoadImage node (at least as of 4f5812b9).
|
||||
"""
|
||||
image_path = _OpenAssetIOHost.instance().resolve_to_path(entity_reference)
|
||||
|
||||
sequence = node_helpers.pillow(Image.open, image_path)
|
||||
|
||||
output_images = []
|
||||
output_masks = []
|
||||
w, h = None, None
|
||||
|
||||
excluded_formats = ["MPO"]
|
||||
|
||||
for frame in ImageSequence.Iterator(sequence):
|
||||
frame = node_helpers.pillow(ImageOps.exif_transpose, frame) # noqa: PLW2901 - frame overwrite
|
||||
|
||||
if frame.mode == "I":
|
||||
frame = frame.point(lambda px: px * (1 / 255)) # noqa: PLW2901 - frame overwrite
|
||||
image = frame.convert("RGB")
|
||||
|
||||
if len(output_images) == 0:
|
||||
w = image.size[0]
|
||||
h = image.size[1]
|
||||
|
||||
if image.size[0] != w or image.size[1] != h:
|
||||
continue
|
||||
|
||||
image = np.array(image).astype(np.float32) / 255.0
|
||||
image = torch.from_numpy(image)[None,]
|
||||
if "A" in frame.getbands():
|
||||
mask = np.array(frame.getchannel("A")).astype(np.float32) / 255.0
|
||||
mask = 1.0 - torch.from_numpy(mask)
|
||||
elif frame.mode == "P" and "transparency" in frame.info:
|
||||
mask = np.array(frame.convert("RGBA").getchannel("A")).astype(np.float32) / 255.0
|
||||
mask = 1.0 - torch.from_numpy(mask)
|
||||
else:
|
||||
mask = torch.zeros((64, 64), dtype=torch.float32, device="cpu")
|
||||
output_images.append(image)
|
||||
output_masks.append(mask.unsqueeze(0))
|
||||
|
||||
if len(output_images) > 1 and sequence.format not in excluded_formats:
|
||||
output_image = torch.cat(output_images, dim=0)
|
||||
output_mask = torch.cat(output_masks, dim=0)
|
||||
else:
|
||||
output_image = output_images[0]
|
||||
output_mask = output_masks[0]
|
||||
|
||||
return output_image, output_mask
|
||||
|
||||
|
||||
# A dictionary that contains all nodes you want to export with their names
|
||||
# NOTE: names should be globally unique
|
||||
NODE_CLASS_MAPPINGS = {"Example": Example}
|
||||
# Plugin registration: node classes.
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"OpenAssetIOResolveImage": ResolveImage,
|
||||
}
|
||||
|
||||
# A dictionary that contains the friendly/humanly readable titles for the nodes
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {"Example": "Example Node"}
|
||||
# Plugin registration: node names.
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"OpenAssetIOResolveImage": "OpenAssetIO Resolve Image",
|
||||
}
|
||||
|
||||
+3
-2
@@ -1,5 +1,6 @@
|
||||
# openassetio-comfyui
|
||||
# Copyright (c) 2025 The Foundry Visionmongers Ltd
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
"""Unit test package for openassetio-comfyui"""
|
||||
"""
|
||||
Unit test package for openassetio-comfyui.
|
||||
"""
|
||||
|
||||
@@ -1,10 +0,0 @@
|
||||
# openassetio-comfyui
|
||||
# Copyright (c) 2025 The Foundry Visionmongers Ltd
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
import os
|
||||
import sys
|
||||
|
||||
# Add the project root directory to Python path
|
||||
# This allows the tests to import the project
|
||||
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
|
||||
@@ -1,8 +0,0 @@
|
||||
# openassetio-comfyui
|
||||
# Copyright (c) 2025 The Foundry Visionmongers Ltd
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
[pytest]
|
||||
testpaths = . # Run tests in the current directory
|
||||
python_files = test_*.py # Run tests in files that start with "test_"
|
||||
norecursedirs = .. # Don't run tests in the parent directory
|
||||
@@ -0,0 +1,27 @@
|
||||
{
|
||||
"entities": {
|
||||
"image_file": {
|
||||
"versions": [
|
||||
{
|
||||
"traits": {
|
||||
"openassetio-mediacreation:usage.Entity": {},
|
||||
"openassetio-mediacreation:twoDimensional.Image": {},
|
||||
"openassetio-mediacreation:content.LocatableContent": {
|
||||
"location": "file://${test_tmp_dir}/cat.png"
|
||||
}
|
||||
}
|
||||
}
|
||||
]
|
||||
},
|
||||
"no_file": {
|
||||
"versions": [
|
||||
{
|
||||
"traits": {
|
||||
"openassetio-mediacreation:usage.Entity": {},
|
||||
"openassetio-mediacreation:twoDimensional.Image": {}
|
||||
}
|
||||
}
|
||||
]
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,10 @@
|
||||
# openassetio-comfyui
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# Copyright 2025 The Foundry Visionmongers Ltd
|
||||
|
||||
[manager]
|
||||
identifier = "org.openassetio.examples.manager.bal"
|
||||
|
||||
[manager.settings]
|
||||
library_path = "${config_dir}/bal_db.json"
|
||||
simulated_query_latency_ms = 0
|
||||
@@ -1,27 +1,218 @@
|
||||
#!/usr/bin/env python
|
||||
# openassetio-comfyui
|
||||
# Copyright (c) 2025 The Foundry Visionmongers Ltd
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
"""
|
||||
Tests for `openassetio-comfyui` package.
|
||||
"""
|
||||
|
||||
"""Tests for `openassetio-comfyui` package."""
|
||||
import contextlib
|
||||
import hashlib
|
||||
# D101: "Missing docstring in public class"
|
||||
# D102: "Missing docstring in public method"
|
||||
# ruff: noqa: D101,D102
|
||||
|
||||
import inspect
|
||||
import pathlib
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
from src.openassetio_comfyui.nodes import Example
|
||||
import torch
|
||||
from PIL import Image
|
||||
|
||||
from openassetio_comfyui.nodes import ResolveImage, _OpenAssetIOHost
|
||||
|
||||
|
||||
class Test_ResolveImage_init:
|
||||
"""
|
||||
Test that the node can be instantiated.
|
||||
"""
|
||||
|
||||
def test_has_correct_type(self, resolve_image_node):
|
||||
assert isinstance(resolve_image_node, ResolveImage)
|
||||
|
||||
|
||||
class Test_ResolveImage_constants:
|
||||
"""
|
||||
Test the node's metadata.
|
||||
"""
|
||||
|
||||
def test_has_expected_values(self):
|
||||
assert ResolveImage.DESCRIPTION == "Resolve images from an asset manager."
|
||||
assert ResolveImage.CATEGORY == "image"
|
||||
assert ResolveImage.FUNCTION == "resolve_image"
|
||||
assert ResolveImage.RETURN_TYPES == ("IMAGE", "MASK")
|
||||
|
||||
def test_function_matches(self):
|
||||
assert inspect.isfunction(getattr(ResolveImage, ResolveImage.FUNCTION))
|
||||
|
||||
|
||||
class Test_ResolveImage_INPUT_TYPES:
|
||||
def test_has_expected_structure(self):
|
||||
input_types = ResolveImage.INPUT_TYPES()
|
||||
|
||||
assert input_types == {
|
||||
"required": {
|
||||
"entity_reference": (
|
||||
"STRING",
|
||||
{
|
||||
"multiline": False,
|
||||
"placeholder": "Entity reference...",
|
||||
"default": "",
|
||||
"tooltip": "The entity to resolve",
|
||||
},
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
class Test_ResolveImage_VALIDATE_INPUTS:
|
||||
def test_when_is_a_reference_then_returns_true(self):
|
||||
valid = ResolveImage.VALIDATE_INPUTS("bal:///")
|
||||
assert valid
|
||||
|
||||
def test_when_not_a_reference_then_returns_false(self):
|
||||
valid = ResolveImage.VALIDATE_INPUTS("notbal:///")
|
||||
assert not valid
|
||||
|
||||
def test_when_no_manager_found_then_raises(self, assert_raises_missing_manager):
|
||||
with assert_raises_missing_manager():
|
||||
ResolveImage.VALIDATE_INPUTS("bal:///")
|
||||
|
||||
|
||||
class Test_ResolveImage_IS_CHANGED:
|
||||
def test_when_reference_resolves_to_file_then_hash_matches_file(self, image_file_hash):
|
||||
actual_file_hash = ResolveImage.IS_CHANGED("bal:///image_file")
|
||||
|
||||
assert actual_file_hash == image_file_hash
|
||||
|
||||
def test_when_reference_has_no_path_then_raises(self, image_file_hash):
|
||||
expected_error_message = (
|
||||
"Failed to resolve entity reference 'bal:///no_file' to a location"
|
||||
)
|
||||
|
||||
with pytest.raises(RuntimeError, match=expected_error_message):
|
||||
ResolveImage.IS_CHANGED("bal:///no_file")
|
||||
|
||||
def test_when_no_manager_found_then_raises(self, assert_raises_missing_manager):
|
||||
with assert_raises_missing_manager():
|
||||
ResolveImage.IS_CHANGED("bal:///image_file")
|
||||
|
||||
|
||||
class Test_ResolveImage_resolve_image:
|
||||
def test_when_reference_resolves_to_file_then_returns_image(
|
||||
self, resolve_image_node, image_file_path
|
||||
):
|
||||
images, masks = resolve_image_node.resolve_image("bal:///image_file")
|
||||
|
||||
assert isinstance(images, torch.Tensor)
|
||||
assert images.shape == (1, 2, 2, 3)
|
||||
assert isinstance(masks, torch.Tensor)
|
||||
assert masks.shape == (1, 2, 2)
|
||||
|
||||
image = images[0]
|
||||
mask = masks[0]
|
||||
|
||||
# Assert channels at each pixel.
|
||||
assert torch.equal(image[0, 0], torch.tensor([0, 0, 0]))
|
||||
assert torch.equal(image[0, 1], torch.tensor([1, 1, 1]))
|
||||
assert torch.equal(image[1, 0], torch.tensor([0, 0, 0]))
|
||||
assert torch.equal(image[1, 1], torch.tensor([1, 1, 1]))
|
||||
|
||||
# Note: mask is `1 - alpha`.
|
||||
assert mask[0, 0].item() == 0
|
||||
assert mask[0, 1].item() == 1
|
||||
assert mask[1, 0].item() == 1
|
||||
assert mask[1, 1].item() == 0
|
||||
|
||||
def test_when_no_manager_found_then_raises(
|
||||
self, resolve_image_node, assert_raises_missing_manager
|
||||
):
|
||||
with assert_raises_missing_manager():
|
||||
resolve_image_node.resolve_image("bal:///image_file")
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def example_node():
|
||||
"""Fixture to create an Example node instance."""
|
||||
return Example()
|
||||
def assert_raises_missing_manager(monkeypatch):
|
||||
"""
|
||||
Fixture to assert that a RuntimeError is raised when no manager can
|
||||
be created.
|
||||
|
||||
Since a lazily-created singleton is used to manage the OpenAssetIO
|
||||
host, this assertion must be tested for every function that uses
|
||||
that singleton.
|
||||
"""
|
||||
|
||||
@contextlib.contextmanager
|
||||
def assert_ctx():
|
||||
monkeypatch.delenv("OPENASSETIO_DEFAULT_CONFIG", raising=False)
|
||||
expected_error_msg = (
|
||||
"Could not create an OpenAssetIO manager instance. Ensure that your OpenAssetIO"
|
||||
" configuration is correct and that the environment variable"
|
||||
" OPENASSETIO_DEFAULT_CONFIG is set to a valid configuration file."
|
||||
)
|
||||
with pytest.raises(RuntimeError, match=expected_error_msg):
|
||||
yield
|
||||
|
||||
return assert_ctx
|
||||
|
||||
|
||||
def test_example_node_initialization(example_node):
|
||||
"""Test that the node can be instantiated."""
|
||||
assert isinstance(example_node, Example)
|
||||
@pytest.fixture
|
||||
def image_file_hash(image_file_path):
|
||||
"""
|
||||
Fixture to provide the hash of the temporary image file.
|
||||
"""
|
||||
hasher = hashlib.sha256()
|
||||
hasher.update(image_file_path.read_bytes())
|
||||
return hasher.hexdigest()
|
||||
|
||||
|
||||
def test_return_types():
|
||||
"""Test the node's metadata."""
|
||||
assert Example.RETURN_TYPES == ("IMAGE",)
|
||||
assert Example.FUNCTION == "test"
|
||||
assert Example.CATEGORY == "Example"
|
||||
@pytest.fixture
|
||||
def image_file_path(monkeypatch, tmp_path):
|
||||
"""
|
||||
Fixture to provide a temporary image file.
|
||||
"""
|
||||
# See bal_db.json - the resolved path will interpolate this env var.
|
||||
monkeypatch.setenv("test_tmp_dir", str(tmp_path))
|
||||
# See bal_db.json - the resolved path for bal:///cat
|
||||
file_path = pathlib.Path(tmp_path) / "cat.png"
|
||||
|
||||
content = np.array(
|
||||
[[[0, 0, 0, 255], [255, 255, 255, 0]], [[0, 0, 0, 0], [255, 255, 255, 255]]]
|
||||
).astype(np.uint8)
|
||||
img = Image.fromarray(content)
|
||||
img.save(file_path)
|
||||
return file_path
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def resolve_image_node():
|
||||
"""
|
||||
Fixture to create an ResolveImage node instance.
|
||||
"""
|
||||
return ResolveImage()
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def openassetio_config_env_var(monkeypatch, resources_dir):
|
||||
"""
|
||||
Fixture to set the OPENASSETIO_DEFAULT_CONFIG environment variable
|
||||
to point to a test config file.
|
||||
"""
|
||||
test_config_path = resources_dir / "openassetio_config.toml"
|
||||
monkeypatch.setenv("OPENASSETIO_DEFAULT_CONFIG", str(test_config_path))
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def reset_singleton_host(monkeypatch):
|
||||
"""
|
||||
Fixture to reset the singleton instance before each test.
|
||||
"""
|
||||
_OpenAssetIOHost._instance = None
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def resources_dir():
|
||||
"""
|
||||
Fixture to provide the path to the resources directory.
|
||||
"""
|
||||
return pathlib.Path(__file__).parent / "resources"
|
||||
|
||||
Reference in New Issue
Block a user