diff --git a/__init__.py b/__init__.py index 716e733..4ea7d6b 100644 --- a/__init__.py +++ b/__init__.py @@ -242,6 +242,7 @@ from .modules.pytorch_wrapper.ptn_linear_model import PtnLinearModel from .modules.pytorch_wrapper.pto_adam import PtoAdam from .modules.pytorch_wrapper.ptv_dataset import PtvDataset from .modules.pytorch_wrapper.ptv_dataset_loader import PtvDatasetLoader +from .modules.pytorch_wrapper.ptv_image_folder_dataset import PtvImageFolderDataset from .modules.pytorch_wrapper.ptv_transforms_to_tensor import PtvTransformsToTensor from .modules.seaborn_wrapper.sns_bar import SNSBar from .modules.seaborn_wrapper.sns_boxplot import SNSBoxplot @@ -492,6 +493,7 @@ NODE_CLASS_MAPPINGS: Dict[str, Type[T]] = { "PtoAdam": PtoAdam, "PtvDataset": PtvDataset, "PtvDatasetLoader": PtvDatasetLoader, + "PtvImageFolderDataset": PtvImageFolderDataset, "PtvTransformsToTensor": PtvTransformsToTensor, "PyDatetimeToString": PyDatetimeToString, "PyFloatToString": PyFloatToString, @@ -751,6 +753,7 @@ NODE_DISPLAY_NAME_MAPPINGS: Dict[str, str] = { "PtoAdam": "Pto Adam", "PtvDataset": "Ptv Dataset", "PtvDatasetLoader": "Ptv Dataset Loader", + "PtvImageFolderDataset": "Ptv Image Folder Dataset", "PtvTransformsToTensor": "Ptv Transforms To Tensor", "PyDatetimeToString": "Py Datetime To String", "PyFloatToString": "Py Float To String", diff --git a/docs/reference/node_reference.md b/docs/reference/node_reference.md index de32f1b..1cf3d9f 100644 --- a/docs/reference/node_reference.md +++ b/docs/reference/node_reference.md @@ -275,6 +275,7 @@ | [Ptn Linear Model](ptn_linear_model.md) | A linear model consisting of dense layers. | | [Ptv Dataset](ptv_dataset.md) | A Torchvision Dataset class wrapper. | | [Ptv Dataset Loader](ptv_dataset_loader.md) | A node to combine the dataset and data loader into a single node. | +| [Ptv Image Folder Dataset](ptv_image_folder_dataset.md) | A Torchvision ImageFolder Dataset class wrapper. | | [Ptv Transforms To Tensor](ptv_transforms_to_tensor.md) | Transforms elements of dataset to PyTorch tensors. | ## PyTorch wrapper - Transform | Node | Description | diff --git a/docs/reference/ptv_image_folder_dataset.md b/docs/reference/ptv_image_folder_dataset.md new file mode 100644 index 0000000..b973b55 --- /dev/null +++ b/docs/reference/ptv_image_folder_dataset.md @@ -0,0 +1,28 @@ +# Ptv Image Folder Dataset +A Torchvision ImageFolder Dataset class wrapper. + +Parameters + +**root** + Specify the root directory that contains subdirectories which of which contains images for each class. + The subdirectory name needs to be named after each class. + +**transform** + Data transforms e.g. transform to PyTorch tensor, normalize image. + Plug in the PTV Transforms node that contains Torchvision Transforms functionality. + +## Input +| Name | Data type | +|---|---| +| root | String | +| transform | Ptvtransform | + +## Output +| Data type | +|---| +| Ptvdataset | + +
+Category: PyTorch wrapper - Training + +ComfyUI Data Analysis Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved. diff --git a/modules/pytorch_wrapper/ptv_image_folder_dataset.py b/modules/pytorch_wrapper/ptv_image_folder_dataset.py new file mode 100644 index 0000000..44ac685 --- /dev/null +++ b/modules/pytorch_wrapper/ptv_image_folder_dataset.py @@ -0,0 +1,73 @@ +from typing import Any, Dict, Optional, Callable, Tuple +import torch +from torchvision.datasets import ImageFolder + +from .utils import get_dataset_full_path + +class PtvImageFolderDataset: + """ + A Torchvision ImageFolder Dataset class wrapper. + + Parameters + + **root** + Specify the root directory that contains subdirectories which of which contains images for each class. + The subdirectory name needs to be named after each class. + + **transform** + Data transforms e.g. transform to PyTorch tensor, normalize image. + Plug in the PTV Transforms node that contains Torchvision Transforms functionality. + + category: PyTorch wrapper - Training + """ + + @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": { + "root": ("STRING", {"default": "", "multiline": False}) + }, + "optional": { + "transform": ("PTVTRANSFORM", {}) + } + } + + RETURN_TYPES: Tuple[str] = ("PTVDATASET",) + FUNCTION: str = "f" + CATEGORY: str = "Training" + + @classmethod + def IS_CHANGED(cls, **kw): + return float("NaN") + + def f(self, + root: str, + transform: Optional[Callable] = None) -> Tuple: + """ + Loads a dataset from Torchvision with specified parameters. + + Args: + root (str): Root directory for storing the dataset. + transform (Optional[Callable]): Transformations to apply to the dataset. + + Returns: + Tuple: A tuple containing the dataset instance. + """ + + with torch.inference_mode(False): + + if root: + dataset_path = get_dataset_full_path(root) + else: # Use default directory + dataset_path = get_dataset_full_path("") + + dc = ImageFolder(root=dataset_path, + transform=transform + ) + return (dc,) diff --git a/pyproject.toml b/pyproject.toml index 8558c08..d73233a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,7 +1,7 @@ [project] name = "data-analysis" description = "Data analysis custom modules for ComfyUI – Integrates Pandas, Seaborn and Matplotlib into ComfyUI, featuring extensive PyTorch wrapper nodes for seamless tensor operations. Experimental PyTorch model training was added in v2.0." -version = "2.0.6" +version = "2.0.7" license = {file = "LICENSE"} [project.urls] diff --git a/tests/pytorch_wrapper/pt_evaluate_classification_conv_model_test2.py b/tests/pytorch_wrapper/pt_evaluate_classification_conv_model_test2.py new file mode 100644 index 0000000..bfb5e7a --- /dev/null +++ b/tests/pytorch_wrapper/pt_evaluate_classification_conv_model_test2.py @@ -0,0 +1,80 @@ +import os +import sys +import unittest +import torch +import torch.nn as nn + +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.ptn_conv_model import PtnConvModel +from pytorch_wrapper.ptv_image_folder_dataset import PtvImageFolderDataset +from pytorch_wrapper.ptv_transforms_to_tensor import PtvTransformsToTensor +from pytorch_wrapper.pt_data_loader import PtDataLoader +from pytorch_wrapper.pt_evaluate_classification_model import PtEvaluateClassificationModel +from pytorch_wrapper.pt_load_model import PtLoadModel + +class TestPtEvaluate(unittest.TestCase): + def setUp(self): + """Set up test instance.""" + # Transform + compose_node = PtvTransformsToTensor() + transform = compose_node.f()[0] + + # Dataset + self.node = PtvImageFolderDataset() + dataset = self.node.f("dog_and_cat/val", + transform=transform)[0] + + # Data loader + data_loader_node = PtDataLoader() + self.data_loader = data_loader_node.f( + dataset=dataset, + batch_size=32, + shuffle=True, + parameters='{"num_workers":1}' + )[0] + + # Model + self.model_node = PtnConvModel() + self.model = self.model_node.f( + input_dim="(3, 512, 512)", + penultimate_dim=0, + output_dim=2, + channel_list="[32,64,128,256,512]", # To 256, 128, 64, 32, 16 + kernel_size_list="[3,3,3,3,3]", + padding_list="[1,1,1,1,1]", + downsample_list="[True,True,True,True,True]" + )[0] + + # Load weights + self.load_model = PtLoadModel() + epochs = 7 + self.loaded_model = self.load_model.f(self.model,f"dog_cat_{epochs}_epochs_conv.pt") + + # Instantiates evaluation node + self.evaluate = PtEvaluateClassificationModel() + + def test_evaluation(self): + """ + Tests evaluation. + + Note that this test requires a trained model file on the file system. + Unit test will fail if you run this for the first time. + Run the unit test for training to generate the weight, and re-run this unit test. + """ + accuracy, precision, recall, f1, num_samples = \ + self.evaluate.f(self.model, self.data_loader, use_gpu=True) + print(f""" + accuracy: {accuracy} + precision: {precision} + recall: {recall} + f1: {f1} + Number of samples: {num_samples} + """) + self.assertTrue(accuracy > 0.7) + self.assertEqual(num_samples, 300) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/pytorch_wrapper/pt_train_classification_conv_model_test2.py b/tests/pytorch_wrapper/pt_train_classification_conv_model_test2.py new file mode 100644 index 0000000..b6fc7d0 --- /dev/null +++ b/tests/pytorch_wrapper/pt_train_classification_conv_model_test2.py @@ -0,0 +1,60 @@ +import os +import sys +import unittest +import torch +import torch.nn as nn + +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.ptn_conv_model import PtnConvModel +from pytorch_wrapper.pto_adam import PtoAdam +from pytorch_wrapper.ptv_image_folder_dataset import PtvImageFolderDataset +from pytorch_wrapper.ptv_transforms_to_tensor import PtvTransformsToTensor +from pytorch_wrapper.pt_data_loader import PtDataLoader +from pytorch_wrapper.pt_train_classification_model import PtTrainClassificationModel +from pytorch_wrapper.pt_save_model import PtSaveModel + + +class TestPtTrainConvModel2(unittest.TestCase): + def setUp(self): + """Set up test instance.""" + compose_node = PtvTransformsToTensor() + transform = compose_node.f()[0] + self.node = PtvImageFolderDataset() + dataset = self.node.f("dog_and_cat/train", + transform=transform)[0] + data_loader_node = PtDataLoader() + self.train_loader = data_loader_node.f( + dataset=dataset, + batch_size=32, + shuffle=True, + parameters='{"num_workers":1}' + )[0] + + self.model_node = PtnConvModel() + self.model = self.model_node.f( + input_dim="(3, 512, 512)", + penultimate_dim=0, + output_dim=2, + channel_list="[32,64,128,256,512]", # To 256, 128, 64, 32, 16 + kernel_size_list="[3,3,3,3,3]", + padding_list="[1,1,1,1,1]", + downsample_list="[True,True,True,True,True]" + )[0] + + self.optimizer = PtoAdam().f(self.model, 0.0001, 0.9, 0.999)[0] + self.trainer = PtTrainClassificationModel() + self.save_model = PtSaveModel() + + def test_1(self): + epochs = 7 + trained_model = self.trainer.f(self.model, self.train_loader, self.optimizer, + epochs, # epochs + use_gpu=True)[0] + self.save_model.f(trained_model, f"dog_cat_{epochs}_epochs_conv.pt") + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/pytorch_wrapper/ptv_image_folder_dataset_test.py b/tests/pytorch_wrapper/ptv_image_folder_dataset_test.py new file mode 100644 index 0000000..bc9de3b --- /dev/null +++ b/tests/pytorch_wrapper/ptv_image_folder_dataset_test.py @@ -0,0 +1,47 @@ +""" +This unit test requires the dog_and_cat dataset to be set up under +project_root/datasets/. +""" +import os +import sys +import unittest + +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.ptv_image_folder_dataset import PtvImageFolderDataset +from pytorch_wrapper.ptv_transforms_to_tensor import PtvTransformsToTensor + + +class TestPtvImageDataset(unittest.TestCase): + def setUp(self): + """Set up test data.""" + self.node = PtvImageFolderDataset() + self.compose_node = PtvTransformsToTensor() + + def test_with_tensor_transform(self): + """Tests downloading with transform""" + transform = self.compose_node.f()[0] + dc = self.node.f("dog_and_cat/train", + transform=transform)[0] + expected_classes = ["cats", "dogs"] + expected_dict = [("cats", 0), ("dogs", 1)] + for i, c in enumerate(dc.classes): + self.assertEqual(c, expected_classes[i]) + + for i, (k, w) in enumerate(dc.class_to_idx.items()): + self.assertEqual(k, expected_dict[i][0]) + self.assertEqual(w, expected_dict[i][1]) + + # Read data + data = next(iter(dc)) + + # Check data shape + self.assertEqual(data[0].size(), (3, 512, 512)) + + # Check class + self.assertTrue(data[1] == 0) # cat + + +if __name__ == "__main__": + unittest.main()