Added Ptv Image Folder Dataset node.

2025-03-06 18:45:32 -08:00
parent be88fd6caa
commit 0d44bd51ec
8 changed files with 293 additions and 1 deletions
@@ -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",
@@ -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 |
@@ -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 |
+
+<HR>
+Category: PyTorch wrapper - Training
+
+ComfyUI Data Analysis Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved.
@@ -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,)
@@ -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]
@@ -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()
@@ -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()
@@ -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()