diff --git a/README.md b/README.md index a37f473..bc68d7f 100644 --- a/README.md +++ b/README.md @@ -43,6 +43,9 @@ Creating sample data and plotting it is incredibly simple, as shown below: ![SNS Line Plot](docs/images/sns_line.png) An example plot +You can also analyze images generated in ComfyUI's image generation workflows. Below screenshot shows a histogram of red channel pixels. +![Red channel pixel distribution](docs/images/red_pixel_distribution.png) + ## Installation Note: GPU is not required to run this custom module. diff --git a/__init__.py b/__init__.py index ddb27b4..9495c18 100644 --- a/__init__.py +++ b/__init__.py @@ -115,14 +115,20 @@ from .modules.pandas_wrapper.pandas_var import PandasVar from .modules.pandas_wrapper.pandas_vertical_concat import PandasVerticalConcat from .modules.pandas_wrapper.pandas_vertical_split import PandasVerticalSplit from .modules.pandas_wrapper.pandas_xs import PandasXs +from .modules.pytorch_wrapper.pt_flatten import PtFlatten from .modules.pytorch_wrapper.pt_float_create import PtFloatCreate from .modules.pytorch_wrapper.pt_from_image import PtFromImage from .modules.pytorch_wrapper.pt_from_latent import PtFromLatent from .modules.pytorch_wrapper.pt_from_numpy import PtFromNumpy from .modules.pytorch_wrapper.pt_int_create import PtIntCreate +from .modules.pytorch_wrapper.pt_size import PtSize +from .modules.pytorch_wrapper.pt_size_create import PtSizeCreate +from .modules.pytorch_wrapper.pt_size_to_numpy import PtSizeToNumpy +from .modules.pytorch_wrapper.pt_size_to_string import PtSizeToString from .modules.pytorch_wrapper.pt_to_image import PtToImage from .modules.pytorch_wrapper.pt_to_latent import PtToLatent from .modules.pytorch_wrapper.pt_to_numpy import PtToNumpy +from .modules.pytorch_wrapper.pt_to_rgb_tensors import PtToRgbTensors from .modules.seaborn_wrapper.sns_bar import SNSBar from .modules.seaborn_wrapper.sns_histogram import SNSHistogram from .modules.seaborn_wrapper.sns_line import SNSLine @@ -248,14 +254,20 @@ NODE_CLASS_MAPPINGS: Dict[str, Type[T]] = { "PandasVerticalConcat": PandasVerticalConcat, "PandasVerticalSplit": PandasVerticalSplit, "PandasXs": PandasXs, + "PtFlatten": PtFlatten, "PtFloatCreate": PtFloatCreate, "PtFromImage": PtFromImage, "PtFromLatent": PtFromLatent, "PtFromNumpy": PtFromNumpy, "PtIntCreate": PtIntCreate, + "PtSize": PtSize, + "PtSizeCreate": PtSizeCreate, + "PtSizeToNumpy": PtSizeToNumpy, + "PtSizeToString": PtSizeToString, "PtToImage": PtToImage, "PtToLatent": PtToLatent, "PtToNumpy": PtToNumpy, + "PtToRgbTensors": PtToRgbTensors, "PyKvFloatCreate": PyKvFloatCreate, "PyKvIntCreate": PyKvIntCreate, "PyKvStringCreate": PyKvStringCreate, @@ -384,14 +396,20 @@ NODE_DISPLAY_NAME_MAPPINGS: Dict[str, str] = { "PandasVerticalConcat": "Pandas Vertical Concat", "PandasVerticalSplit": "Pandas Vertical Split", "PandasXs": "Pandas Xs", + "PtFlatten": "Pt Flatten", "PtFloatCreate": "Pt Float Create", "PtFromImage": "Pt From Image", "PtFromLatent": "Pt From Latent", "PtFromNumpy": "Pt From Numpy", "PtIntCreate": "Pt Int Create", + "PtSize": "Pt Size", + "PtSizeCreate": "Pt Size Create", + "PtSizeToNumpy": "Pt Size To Numpy", + "PtSizeToString": "Pt Size To String", "PtToImage": "Pt To Image", "PtToLatent": "Pt To Latent", "PtToNumpy": "Pt To Numpy", + "PtToRgbTensors": "Pt To Rgb Tensors", "PyKvFloatCreate": "Py Kv Float Create", "PyKvIntCreate": "Py Kv Int Create", "PyKvStringCreate": "Py Kv String Create", diff --git a/docs/images/red_pixel_distribution.png b/docs/images/red_pixel_distribution.png new file mode 100644 index 0000000..ceb12ce Binary files /dev/null and b/docs/images/red_pixel_distribution.png differ diff --git a/docs/reference/node_reference.md b/docs/reference/node_reference.md index 60b29f9..4812790 100644 --- a/docs/reference/node_reference.md +++ b/docs/reference/node_reference.md @@ -113,6 +113,11 @@ | [Pandas Log](pandas_log.md) | Applying the natural logarithm function to a pandas DataFrame and converting non-numeric values to NaN. | | [Pandas Sin](pandas_sin.md) | Applies the sine function to a pandas DataFrame and converting non-numeric values to NaN. | | [Pandas Tan](pandas_tan.md) | Applies the tangent function to a pandas DataFrame and converting non-numeric values to NaN. | +## Miscellaneous +| Node | Description | +| --- | --- | +| [Pt Flatten](pt_flatten.md) | Flattens a PyTorch tensor into a 1D tensor. | +| [Pt To Rgb Tensors](pt_to_rgb_tensors.md) | Splits a PyTorch tensor into R, G, and B tensors. | ## Numpy | Node | Description | | --- | --- | @@ -136,6 +141,10 @@ | [Pt From Latent](pt_from_latent.md) | Casts a latent tensor as a PyTorch tensor. | | [Pt From Numpy](pt_from_numpy.md) | Converts a NumPy ndarray to a PyTorch tensor while preserving its data type. | | [Pt Int Create](pt_int_create.md) | Creates a PyTorch tensor with 32-bit integer | +| [Pt Size](pt_size.md) | Extracts the PyTorch Size object of a PyTorch tensor using the size() method. | +| [Pt Size Create](pt_size_create.md) | Creates a PyTorch Size using values entered in the text field. | +| [Pt Size To Numpy](pt_size_to_numpy.md) | Converts PyTorch Size object to NumPy ndarray. | +| [Pt Size To String](pt_size_to_string.md) | Converts PyTorch Size object to a Python string. | | [Pt To Image](pt_to_image.md) | Casts a PyTorch tensor as an Image tensor. | | [Pt To Latent](pt_to_latent.md) | Casts a PyTorch tensor as a latent tensor. | | [Pt To Numpy](pt_to_numpy.md) | Converts PyTorch tensor to NumPy ndarray. | diff --git a/docs/reference/pt_flatten.md b/docs/reference/pt_flatten.md new file mode 100644 index 0000000..c9f79d6 --- /dev/null +++ b/docs/reference/pt_flatten.md @@ -0,0 +1,19 @@ +# Pt Flatten +Flattens a PyTorch tensor into a 1D tensor. + +Category: PyTorch wrapper + +## Input +| Name | Data type | +|---|---| +| tens | Tensor | + +## Output +| Data type | +|---| +| Tensor | + +
+Category: Miscellaneous + +ComfyUI Data Analysis Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved. diff --git a/docs/reference/pt_size.md b/docs/reference/pt_size.md new file mode 100644 index 0000000..7a4bab0 --- /dev/null +++ b/docs/reference/pt_size.md @@ -0,0 +1,17 @@ +# Pt Size +Extracts the PyTorch Size object of a PyTorch tensor using the size() method. + +## Input +| Name | Data type | +|---|---| +| tens | Tensor | + +## Output +| Data type | +|---| +| Ptsize | + +
+Category: PyTorch wrapper + +ComfyUI Data Analysis Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved. diff --git a/docs/reference/pt_size_create.md b/docs/reference/pt_size_create.md new file mode 100644 index 0000000..256505f --- /dev/null +++ b/docs/reference/pt_size_create.md @@ -0,0 +1,17 @@ +# Pt Size Create +Creates a PyTorch Size using values entered in the text field. + +## Input +| Name | Data type | +|---|---| +| data | String | + +## Output +| Data type | +|---| +| Ptsize | + +
+Category: PyTorch wrapper + +ComfyUI Data Analysis Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved. diff --git a/docs/reference/pt_size_to_numpy.md b/docs/reference/pt_size_to_numpy.md new file mode 100644 index 0000000..853ebff --- /dev/null +++ b/docs/reference/pt_size_to_numpy.md @@ -0,0 +1,17 @@ +# Pt Size To Numpy +Converts PyTorch Size object to NumPy ndarray. + +## Input +| Name | Data type | +|---|---| +| sz | Ptsize | + +## Output +| Data type | +|---| +| Ndarray | + +
+Category: PyTorch wrapper + +ComfyUI Data Analysis Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved. diff --git a/docs/reference/pt_size_to_string.md b/docs/reference/pt_size_to_string.md new file mode 100644 index 0000000..52f3ec0 --- /dev/null +++ b/docs/reference/pt_size_to_string.md @@ -0,0 +1,17 @@ +# Pt Size To String +Converts PyTorch Size object to a Python string. + +## Input +| Name | Data type | +|---|---| +| sz | Ptsize | + +## Output +| Data type | +|---| +| String | + +
+Category: PyTorch wrapper + +ComfyUI Data Analysis Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved. diff --git a/docs/reference/pt_to_rgb_tensors.md b/docs/reference/pt_to_rgb_tensors.md new file mode 100644 index 0000000..efd507a --- /dev/null +++ b/docs/reference/pt_to_rgb_tensors.md @@ -0,0 +1,21 @@ +# Pt To Rgb Tensors +Splits a PyTorch tensor into R, G, and B tensors. + +Category: PyTorch wrapper + +## Input +| Name | Data type | +|---|---| +| tens | Tensor | + +## Output +| Data type | +|---| +| Tensor | +| Tensor | +| Tensor | + +
+Category: Miscellaneous + +ComfyUI Data Analysis Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved. diff --git a/modules/pytorch_wrapper/pt_flatten.py b/modules/pytorch_wrapper/pt_flatten.py new file mode 100644 index 0000000..791062d --- /dev/null +++ b/modules/pytorch_wrapper/pt_flatten.py @@ -0,0 +1,40 @@ +from typing import Any, Dict +import torch + + +class PtFlatten: + """ + Flattens a PyTorch tensor into a 1D tensor. + + Category: PyTorch wrapper + """ + + @classmethod + def INPUT_TYPES(cls) -> Dict[str, Any]: + """ + Defines the required input types. + + Returns: + Dict[str, Any]: A dictionary of required input types. + """ + return { + "required": { + "tens": ("TENSOR", {}), + } + } + + RETURN_TYPES: tuple = ("TENSOR",) + FUNCTION: str = "f" + CATEGORY: str = "PyTorch wrapper" + + def f(self, tens: torch.Tensor) -> tuple: + """ + Flattens a PyTorch tensor into a 1D tensor. + + Args: + tens (torch.Tensor): A PyTorch tensor of any shape. + + Returns: + tuple: A tuple containing the flattened tensor. + """ + return (tens.flatten(),) diff --git a/modules/pytorch_wrapper/pt_from_latent.py b/modules/pytorch_wrapper/pt_from_latent.py index fc313e7..6bc9ddc 100644 --- a/modules/pytorch_wrapper/pt_from_latent.py +++ b/modules/pytorch_wrapper/pt_from_latent.py @@ -37,4 +37,4 @@ class PtFromLatent: Returns: tuple: A tuple containing the PyTorch tensor. """ - return (latent,) + return (latent["samples"],) diff --git a/modules/pytorch_wrapper/pt_size.py b/modules/pytorch_wrapper/pt_size.py new file mode 100644 index 0000000..6965cd4 --- /dev/null +++ b/modules/pytorch_wrapper/pt_size.py @@ -0,0 +1,41 @@ +from typing import Any, Dict +import torch +import ast + + +class PtSize: + """ + Extracts the PyTorch Size object of a PyTorch tensor using the size() method. + + category: PyTorch wrapper + """ + + @classmethod + def INPUT_TYPES(cls) -> Dict[str, Any]: + """ + Defines the input types for the function. + + Returns: + Dict[str, Any]: A dictionary specifying required input types. + """ + return { + "required": { + "tens": ("TENSOR",), + } + } + + RETURN_TYPES: tuple = ("PTSIZE",) + FUNCTION: str = "f" + CATEGORY: str = "PyTorch wrapper" + + def f(self, tens: torch.Tensor) -> tuple: + """ + Extracts the PyTorch Size object of a PyTorch tensor using the size() method. + + Args: + tens (torch.Tensor): The tensor. + + Returns: + tuple: A tuple containing the PyTorch Size. + """ + return (tens.size(),) diff --git a/modules/pytorch_wrapper/pt_size_create.py b/modules/pytorch_wrapper/pt_size_create.py new file mode 100644 index 0000000..9403056 --- /dev/null +++ b/modules/pytorch_wrapper/pt_size_create.py @@ -0,0 +1,46 @@ +from typing import Any, Dict +import torch +import ast + + +class PtSizeCreate: + """ + Creates a PyTorch Size using values entered in the text field. + + category: PyTorch wrapper + """ + + @classmethod + def INPUT_TYPES(cls) -> Dict[str, Any]: + """ + Defines the input types for the function. + + Returns: + Dict[str, Any]: A dictionary specifying required input types. + """ + return { + "required": { + "data": ("STRING", {"multiline": True}), + } + } + + RETURN_TYPES: tuple = ("PTSIZE",) + FUNCTION: str = "f" + CATEGORY: str = "PyTorch wrapper" + + def f(self, data: str) -> tuple: + """ + Creates a PyTorch Size using values entered in the text field. + + Args: + data (str): String value in the text field. + + Returns: + tuple: A tuple containing the PyTorch Size. + """ + # Convert string to Python list + list_data = ast.literal_eval(data) + + # Convert list to PyTorch Size + size = torch.Size(list_data) + return (size,) diff --git a/modules/pytorch_wrapper/pt_size_to_numpy.py b/modules/pytorch_wrapper/pt_size_to_numpy.py new file mode 100644 index 0000000..e0f3900 --- /dev/null +++ b/modules/pytorch_wrapper/pt_size_to_numpy.py @@ -0,0 +1,41 @@ +from typing import Any, Dict +import torch +import numpy as np + +class PtSizeToNumpy: + """ + Converts PyTorch Size object to NumPy ndarray. + + category: PyTorch wrapper + """ + + @classmethod + def INPUT_TYPES(cls) -> Dict[str, Any]: + """ + Defines the input types for the function. + + Returns: + Dict[str, Any]: A dictionary specifying required input types. + """ + return { + "required": { + "sz": ("PTSIZE", {}), + } + } + + RETURN_TYPES: tuple = ("NDARRAY",) + FUNCTION: str = "f" + CATEGORY: str = "PyTorch wrapper" + + def f(self, sz: torch.Size) -> tuple: + """ + Converts PyTorch Size object to NumPy ndarray. + + Args: + sz (torch.Size): PyTorch Size + + Returns: + tuple: A tuple containing the ndarray. + """ + array = np.array(sz, dtype=int) + return (array,) diff --git a/modules/pytorch_wrapper/pt_size_to_string.py b/modules/pytorch_wrapper/pt_size_to_string.py new file mode 100644 index 0000000..25f7bed --- /dev/null +++ b/modules/pytorch_wrapper/pt_size_to_string.py @@ -0,0 +1,40 @@ +from typing import Any, Dict +import torch + + +class PtSizeToString: + """ + Converts PyTorch Size object to a Python string. + + category: PyTorch wrapper + """ + + @classmethod + def INPUT_TYPES(cls) -> Dict[str, Any]: + """ + Defines the input types for the function. + + Returns: + Dict[str, Any]: A dictionary specifying required input types. + """ + return { + "required": { + "sz": ("PTSIZE", {}), + } + } + + RETURN_TYPES: tuple = ("STRING",) + FUNCTION: str = "f" + CATEGORY: str = "PyTorch wrapper" + + def f(self, sz: torch.Size) -> tuple: + """ + Converts PyTorch Size object to a Python string. + + Args: + sz (torch.Size): PyTorch Size + + Returns: + tuple: A tuple containing the string. + """ + return (str(sz),) diff --git a/modules/pytorch_wrapper/pt_to_latent.py b/modules/pytorch_wrapper/pt_to_latent.py index abee2b9..240d3cc 100644 --- a/modules/pytorch_wrapper/pt_to_latent.py +++ b/modules/pytorch_wrapper/pt_to_latent.py @@ -37,4 +37,4 @@ class PtToLatent: Returns: tuple: A tuple containing the latent tensor. """ - return (tens,) + return ({"samples":tens},) diff --git a/modules/pytorch_wrapper/pt_to_rgb_tensors.py b/modules/pytorch_wrapper/pt_to_rgb_tensors.py new file mode 100644 index 0000000..66f6a7c --- /dev/null +++ b/modules/pytorch_wrapper/pt_to_rgb_tensors.py @@ -0,0 +1,50 @@ +from typing import Any, Dict +import torch + + +class PtToRgbTensors: + """ + Splits a PyTorch tensor into R, G, and B tensors. + + Category: PyTorch wrapper + """ + + @classmethod + def INPUT_TYPES(cls) -> Dict[str, Any]: + """ + Defines the required input types. + + Returns: + Dict[str, Any]: A dictionary of required input types. + """ + return { + "required": { + "tens": ("TENSOR", {}), + } + } + + RETURN_NAMES: tuple = ("R", "G", "B") + RETURN_TYPES: tuple = ("TENSOR", "TENSOR", "TENSOR") + FUNCTION: str = "f" + CATEGORY: str = "PyTorch wrapper" + + def f(self, tens: torch.Tensor) -> tuple: + """ + Splits a PyTorch tensor into R, G, and B tensors. + + Args: + tens (torch.Tensor): A PyTorch tensor of shape [B, H, W, C] or [H, W, C]. + Assumes the last axis represents the color channels. + + Returns: + tuple: A tuple containing the R, G, and B tensors. + If the input tensor has a batch axis, it will be retained. + + Raises: + ValueError: If the input tensor does not have exactly 3 channels. + """ + if tens.size()[-1] != 3: + raise ValueError(f"Expected a 3-channel tensor, but got {tens.size()[-1]} channels.") + + r, g, b = tens[..., 0], tens[..., 1], tens[..., 2] + return (r, g, b) diff --git a/pyproject.toml b/pyproject.toml index d32b204..5bbb3dd 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,7 +1,7 @@ [project] name = "data-analysis" description = "Data analysis custom modules for ComfyUI - Use Pandas & Matplotlib within ComfyUI" -version = "1.3.2" +version = "1.3.3" license = {file = "LICENSE"} [project.urls] diff --git a/tests/pytorch_wrapper/pt_flatten_test.py b/tests/pytorch_wrapper/pt_flatten_test.py new file mode 100644 index 0000000..5a66695 --- /dev/null +++ b/tests/pytorch_wrapper/pt_flatten_test.py @@ -0,0 +1,39 @@ +import os +import sys +import unittest +import torch + +PROJECT_ROOT = os.path.realpath(os.path.join(os.path.dirname(__file__), "..", "..")) +MODULE_ROOT = os.path.join(PROJECT_ROOT, "modules") +sys.path.append(MODULE_ROOT) + +from pytorch_wrapper.pt_flatten import PtFlatten + + +class TestPtFlatten(unittest.TestCase): + def setUp(self): + """Set up test instance.""" + self.node = PtFlatten() + + # Test cases: (Input tensor, Expected flattened tensor shape) + self.test_cases = [ + (torch.randn(3, 4, 5), (60,)), # 3D Tensor + (torch.randn(2, 3, 4, 5), (120,)), # 4D Tensor + (torch.randn(10, 10), (100,)), # 2D Tensor + (torch.randn(1), (1,)), # 1D Tensor remains unchanged + ] + + def test_flatten(self): + """Test flattening tensors correctly.""" + for tensor, expected_shape in self.test_cases: + with self.subTest(tensor_shape=tensor.shape): + flattened, = self.node.f(tensor) + + # Ensure return value is a tensor + self.assertTrue(isinstance(flattened, torch.Tensor), "Output is not a tensor") + + # Ensure shape matches expected flattened shape + self.assertEqual(flattened.shape, expected_shape, f"Shape mismatch: expected {expected_shape}, got {flattened.shape}") + +if __name__ == "__main__": + unittest.main() diff --git a/tests/pytorch_wrapper/pt_size_create_test.py b/tests/pytorch_wrapper/pt_size_create_test.py new file mode 100644 index 0000000..a229ac8 --- /dev/null +++ b/tests/pytorch_wrapper/pt_size_create_test.py @@ -0,0 +1,46 @@ +import os +import sys +import unittest +import torch + +PROJECT_ROOT = os.path.realpath(os.path.join(os.path.dirname(__file__), "..", "..")) +MODULE_ROOT = os.path.join(PROJECT_ROOT, "modules") +sys.path.append(MODULE_ROOT) + +from pytorch_wrapper.pt_size_create import PtSizeCreate + +class TestPtSizeCreate(unittest.TestCase): + def setUp(self): + """Set up test data.""" + self.node = PtSizeCreate() + + # Test cases with torch.Size objects + self.test_cases = [ + # Empty size (zero-rank tensor shape) + ("[]", torch.Size([])), + # Rank 1 (single-dimension size) + ("[3]", torch.Size([3])), + # Rank 2 (matrix size) + ("[3, 4]", torch.Size([3, 4])), + # Rank 3 (3D tensor size) + ("[2, 3, 4]", torch.Size([2, 3, 4])), + # Rank 4 (higher-dimensional tensor size) + ("[1, 2, 3, 4]", torch.Size([1, 2, 3, 4])), + # Rank 5 (very high-dimensional tensor size) + ("[5, 4, 3, 2, 1]", torch.Size([5, 4, 3, 2, 1])), + ] + + def test_create_size(self): + """Test PyTorch Size creation from various structured inputs.""" + for data_str, expected_size in self.test_cases: + with self.subTest(data=data_str): + retval = self.node.f(data_str)[0] + + # Ensure the returned value is a PyTorch Size + self.assertTrue(isinstance(retval, torch.Size), f"Expected Size, got {type(retval)}") + + # Ensure the value matches expected size + self.assertEqual(retval, expected_size, f"Value mismatch for input: {data_str}") + +if __name__ == "__main__": + unittest.main() diff --git a/tests/pytorch_wrapper/pt_size_test.py b/tests/pytorch_wrapper/pt_size_test.py new file mode 100644 index 0000000..c268f48 --- /dev/null +++ b/tests/pytorch_wrapper/pt_size_test.py @@ -0,0 +1,44 @@ +import os +import sys +import unittest +import torch + +PROJECT_ROOT = os.path.realpath(os.path.join(os.path.dirname(__file__), "..", "..")) +MODULE_ROOT = os.path.join(PROJECT_ROOT, "modules") +sys.path.append(MODULE_ROOT) + +from pytorch_wrapper.pt_size import PtSize + +class TestPtSize(unittest.TestCase): + def setUp(self): + """Set up test data.""" + self.node = PtSize() + + # Test cases with torch.Tensor objects + self.test_cases = [ + # Rank 0 (scalar tensor) + (torch.tensor(42), torch.Size([])), + # Rank 1 (vector tensor) + (torch.tensor([1, 2, 3]), torch.Size([3])), + # Rank 2 (matrix tensor) + (torch.tensor([[1, 2], [3, 4]]), torch.Size([2, 2])), + # Rank 3 (3D tensor) + (torch.tensor([[[1], [2]], [[3], [4]]]), torch.Size([2, 2, 1])), + # Rank 4 (4D tensor) + (torch.tensor([[[[1, 2], [3, 4]]]]), torch.Size([1, 1, 2, 2])), + ] + + def test_extract_size(self): + """Test PyTorch Size extraction from tensors.""" + for tensor, expected_size in self.test_cases: + with self.subTest(data=tensor): + retval = self.node.f(tensor)[0] + + # Ensure the returned value is a PyTorch Size + self.assertTrue(isinstance(retval, torch.Size), f"Expected Size, got {type(retval)}") + + # Ensure the value matches expected size + self.assertEqual(retval, expected_size, f"Value mismatch for input: {tensor}") + +if __name__ == "__main__": + unittest.main() diff --git a/tests/pytorch_wrapper/pt_size_to_numpy_test.py b/tests/pytorch_wrapper/pt_size_to_numpy_test.py new file mode 100644 index 0000000..feca5ef --- /dev/null +++ b/tests/pytorch_wrapper/pt_size_to_numpy_test.py @@ -0,0 +1,55 @@ +import os +import sys +import unittest +import torch +import numpy as np + +PROJECT_ROOT = os.path.realpath(os.path.join(os.path.dirname(__file__), "..", "..")) +MODULE_ROOT = os.path.join(PROJECT_ROOT, "modules") +sys.path.append(MODULE_ROOT) + +from pytorch_wrapper.pt_size_to_numpy import PtSizeToNumpy + +class TestPtSizeToNumpy(unittest.TestCase): + def setUp(self): + """Set up test data.""" + self.node = PtSizeToNumpy() + + # Test cases with torch.Size objects + self.test_cases = [ + # Empty size (zero-rank tensor shape) + (torch.Size([]), np.array([], dtype=int)), + # Rank 1 (single-dimension size) + (torch.Size([3]), np.array([3], dtype=int)), + # Rank 2 (matrix size) + (torch.Size([3, 4]), np.array([3, 4], dtype=int)), + # Rank 3 (3D tensor size) + (torch.Size([2, 3, 4]), np.array([2, 3, 4], dtype=int)), + # Rank 4 (higher-dimensional tensor size) + (torch.Size([1, 2, 3, 4]), np.array([1, 2, 3, 4], dtype=int)), + # Rank 5 (very high-dimensional tensor size) + (torch.Size([5, 4, 3, 2, 1]), np.array([5, 4, 3, 2, 1], dtype=int)), + ] + + def test_convert_size_to_numpy(self): + """Test conversion of PyTorch Size to NumPy ndarray.""" + for torch_size, expected_array in self.test_cases: + with self.subTest(data=torch_size): + retval = self.node.f(torch_size)[0] + + # Ensure the returned value is a NumPy array + self.assertTrue(isinstance(retval, np.ndarray), f"Expected ndarray, got {type(retval)}") + + # Ensure the dtype is integer + self.assertEqual(retval.dtype, expected_array.dtype, f"Dtype mismatch for input: {torch_size}") + + # Ensure the shape matches + self.assertEqual(retval.shape, expected_array.shape, + f"Shape mismatch for input: {torch_size} (Expected {expected_array.shape}, got {retval.shape})") + + # Ensure values match + np.testing.assert_array_equal(retval, expected_array, + f"Value mismatch for input: {torch_size}") + +if __name__ == "__main__": + unittest.main() diff --git a/tests/pytorch_wrapper/pt_size_to_string_test.py b/tests/pytorch_wrapper/pt_size_to_string_test.py new file mode 100644 index 0000000..69f9b54 --- /dev/null +++ b/tests/pytorch_wrapper/pt_size_to_string_test.py @@ -0,0 +1,46 @@ +import os +import sys +import unittest +import torch + +PROJECT_ROOT = os.path.realpath(os.path.join(os.path.dirname(__file__), "..", "..")) +MODULE_ROOT = os.path.join(PROJECT_ROOT, "modules") +sys.path.append(MODULE_ROOT) + +from pytorch_wrapper.pt_size_to_string import PtSizeToString + +class TestPtSizeToString(unittest.TestCase): + def setUp(self): + """Set up test data.""" + self.node = PtSizeToString() + + # Test cases with torch.Size objects + self.test_cases = [ + # Empty size (zero-rank tensor shape) + (torch.Size([]), "torch.Size([])"), + # Rank 1 (single-dimension size) + (torch.Size([3]), "torch.Size([3])"), + # Rank 2 (matrix size) + (torch.Size([3, 4]), "torch.Size([3, 4])"), + # Rank 3 (3D tensor size) + (torch.Size([2, 3, 4]), "torch.Size([2, 3, 4])"), + # Rank 4 (higher-dimensional tensor size) + (torch.Size([1, 2, 3, 4]), "torch.Size([1, 2, 3, 4])"), + # Rank 5 (very high-dimensional tensor size) + (torch.Size([5, 4, 3, 2, 1]), "torch.Size([5, 4, 3, 2, 1])"), + ] + + def test_convert_size_to_string(self): + """Test conversion of PyTorch Size to string.""" + for torch_size, expected_string in self.test_cases: + with self.subTest(data=torch_size): + retval = self.node.f(torch_size)[0] + + # Ensure the returned value is a string + self.assertTrue(isinstance(retval, str), f"Expected str, got {type(retval)}") + + # Ensure the value matches expected string + self.assertEqual(retval, expected_string, f"Value mismatch for input: {torch_size}") + +if __name__ == "__main__": + unittest.main() diff --git a/tests/pytorch_wrapper/pt_to_rgb_tensors_test.py b/tests/pytorch_wrapper/pt_to_rgb_tensors_test.py new file mode 100644 index 0000000..69f411a --- /dev/null +++ b/tests/pytorch_wrapper/pt_to_rgb_tensors_test.py @@ -0,0 +1,69 @@ +import os +import sys +import unittest +import torch + +PROJECT_ROOT = os.path.realpath(os.path.join(os.path.dirname(__file__), "..", "..")) +MODULE_ROOT = os.path.join(PROJECT_ROOT, "modules") +sys.path.append(MODULE_ROOT) + +from pytorch_wrapper.pt_to_rgb_tensors import PtToRgbTensors + + +class TestPtToRgbTensors(unittest.TestCase): + def setUp(self): + """Set up test instance.""" + self.node = PtToRgbTensors() + + # Test cases: (Input tensor, Expected R, Expected G, Expected B) + self.test_cases = [ + # 3D Tensor: [H, W, C] + ( + torch.tensor([[[255, 128, 64], [0, 127, 255]]], dtype=torch.uint8), + torch.tensor([[255, 0]], dtype=torch.uint8), + torch.tensor([[128, 127]], dtype=torch.uint8), + torch.tensor([[64, 255]], dtype=torch.uint8), + ), + + # 4D Tensor: [B, H, W, C] + ( + torch.tensor( + [[[[255, 0, 0], [0, 255, 0]], [[0, 0, 255], [255, 255, 255]]]], dtype=torch.uint8 + ), + torch.tensor([[[255, 0], [0, 255]]], dtype=torch.uint8), + torch.tensor([[[0, 255], [0, 255]]], dtype=torch.uint8), + torch.tensor([[[0, 0], [255, 255]]], dtype=torch.uint8), + ), + ] + + def test_split_rgb(self): + """Test splitting RGB tensors correctly.""" + for tensor, expected_r, expected_g, expected_b in self.test_cases: + with self.subTest(tensor_shape=tensor.shape): + r, g, b = self.node.f(tensor) + + # Ensure return values are tensors + self.assertTrue(isinstance(r, torch.Tensor), "R is not a tensor") + self.assertTrue(isinstance(g, torch.Tensor), "G is not a tensor") + self.assertTrue(isinstance(b, torch.Tensor), "B is not a tensor") + + # Ensure shapes match + self.assertEqual(r.shape, expected_r.shape, f"R shape mismatch: expected {expected_r.shape}, got {r.shape}") + self.assertEqual(g.shape, expected_g.shape, f"G shape mismatch: expected {expected_g.shape}, got {g.shape}") + self.assertEqual(b.shape, expected_b.shape, f"B shape mismatch: expected {expected_b.shape}, got {b.shape}") + + # Ensure values match + torch.testing.assert_close(r, expected_r, msg="R values mismatch") + torch.testing.assert_close(g, expected_g, msg="G values mismatch") + torch.testing.assert_close(b, expected_b, msg="B values mismatch") + + def test_invalid_input(self): + """Test error handling for non-3-channel input.""" + invalid_tensor = torch.rand(1, 512, 512, 4) # 4 channels instead of 3 + with self.assertRaises(ValueError) as context: + self.node.f(invalid_tensor) + self.assertIn("Expected a 3-channel tensor", str(context.exception)) + + +if __name__ == "__main__": + unittest.main()