Added data augmentation transform and CIFAR-10 workflows.

README.md

@@ -12,6 +12,12 @@ Check out [this tutorial on classifying dogs and cats](docs/dog_cat_classificati
 
 For visualization or analysis features, install ComfyUI-Data-Analysis as well.
 
+The workflow directory contains an easy-to-use workflow for training a high-performing ResNet model on CIFAR-10.
+
+**ResNet model**
+![Train](docs/images/resnet_train.png)
+![Eval](docs/images/resnet_eval.png)
+
 # Getting Started
 If you're interested in direct operations on PyTorch tensors, check out the [Node Reference](docs/reference/node_reference.md) to explore the available nodes and see what operations are supported.
 

__init__.py

@@ -106,6 +106,7 @@ from .modules.pytorch_wrapper.ptv_dataset import PtvDataset
 from .modules.pytorch_wrapper.ptv_dataset_len import PtvDatasetLen
 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_data_augment import PtvTransformsDataAugment
 from .modules.pytorch_wrapper.ptv_transforms_resize import PtvTransformsResize
 from .modules.pytorch_wrapper.ptv_transforms_to_tensor import PtvTransformsToTensor
 T = TypeVar("T")
@@ -223,6 +224,7 @@ NODE_CLASS_MAPPINGS: Dict[str, Type[T]] = {
     "PtvDatasetLen": PtvDatasetLen,
     "PtvDatasetLoader": PtvDatasetLoader,
     "PtvImageFolderDataset": PtvImageFolderDataset,
+    "PtvTransformsDataAugment": PtvTransformsDataAugment,
     "PtvTransformsResize": PtvTransformsResize,
     "PtvTransformsToTensor": PtvTransformsToTensor,
 }
@@ -340,6 +342,7 @@ NODE_DISPLAY_NAME_MAPPINGS: Dict[str, str] = {
     "PtvDatasetLen": "Ptv Dataset Len",
     "PtvDatasetLoader": "Ptv Dataset Loader",
     "PtvImageFolderDataset": "Ptv Image Folder Dataset",
+    "PtvTransformsDataAugment": "Ptv Transforms Data Augment",
     "PtvTransformsResize": "Ptv Transforms Resize",
     "PtvTransformsToTensor": "Ptv Transforms To Tensor",
 }

docs/model_training.md

@@ -21,20 +21,24 @@ Check out [the dog and cat classification tutorial](dog_cat_classification_model
 
 ## Public Datasets
 
-Currently, only Fashion MNIST and CIFAR10 have been tested. Other public datasets have not been verified yet, though they may work.  
+Currently, only Fashion MNIST and CIFAR-10 have been tested. Other public datasets have not been verified yet, though they may work.  
 If you try one and it doesn't work, please create an issue so I can investigate adding support.
 
-To use the Fashion MNIST dataset, simply drag and drop the workflow files below into ComfyUI. The model will be downloaded automatically.
+To use these datasets, simply drag and drop the workflow files below into ComfyUI. The model will be downloaded automatically.  Make sure that you update the directory to save the trained model.
 
 Linear model  
-* Training: examples/fashion_mnist_train.json  
-* Evaluation: examples/fashion_mnist_eval.json
+* Training: examples/workflows/fashion_mnist_train.json  
+* Evaluation: examples/workflows/fashion_mnist_eval.json
 
 Convolutional model  
-* Training: examples/fashion_mnist_train_conv.json  
-* Evaluation: examples/fashion_mnist_eval_conv.json
+* Training: examples/workflows/fashion_mnist_train_conv.json  
+* Evaluation: examples/workflows/fashion_mnist_eval_conv.json
 
-These training workflows train a FashionMNIST model from scratch effortlessly.
+ResNet model
+* Training: examples/workflows/cifar10_train_v8.json
+* Evaluation: examples/workflows/cifar10_eval_v8.json
+
+This model achieved 93.12% accuracy on CIFAR-10. While not state-of-the-art, it is a strong result. Due to random initialization, data augmentation, and dataset shuffling, accuracy may vary slightly between runs. If your accuracy is significantly lower, consider running it multiple times or checking for potential issues in the training setup.
 
 **Note: Before you try these workflows, please install ComfyUI-Data-Analysis extension.**
 
@@ -48,9 +52,15 @@ These training workflows train a FashionMNIST model from scratch effortlessly.
 *Close up of the conv node*
 ![Train](images/conv_train2.png)
 
+**ResNet model**
+![Train](images/resnet_train.png)
+
 ### Eval workflow
 **Linear model**
 ![Eval](images/fashion_mnist_eval.png)
 
 **Conv model**
 ![Eval](images/conv_eval.png)
+
+**ResNet model**
+![Eval](images/resnet_eval.png)

docs/reference/node_reference.md

@@ -144,6 +144,7 @@
 | [Ptv Dataset Len](ptv_dataset_len.md) | A wrapper class that calls Python len on a dataset. |
 | [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 Data Augment](ptv_transforms_data_augment.md) | Applies data augmentation transformations to dataset elements. |
 | [Ptv Transforms Resize](ptv_transforms_resize.md) | Resizes and transforms elements of dataset to PyTorch tensors. |
 | [Ptv Transforms To Tensor](ptv_transforms_to_tensor.md) | Transforms elements of dataset to PyTorch tensors. |
 ## PyTorch wrapper - Transform

docs/reference/ptv_transforms_data_augment.md

@@ -0,0 +1,24 @@
+# Ptv Transforms Data Augment
+Applies data augmentation transformations to dataset elements.
+Supported transformations include flipping, rotation, translation, and scaling.
+
+## Input
+| Name | Data type |
+|---|---|
+| h_flip_prob | Float |
+| v_flip_prob | Float |
+| rotate_degree | Float |
+| h_translate_ratio | Float |
+| v_translate_ratio | Float |
+| min_scale | Float |
+| max_scale | Float |
+
+## Output
+| Data type |
+|---|
+| Ptvtransform |
+
+<HR>
+Category: PyTorch wrapper - Training
+
+ComfyUI Data Analysis Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved.

modules/pytorch_wrapper/ptv_transforms_data_augment.py

@@ -0,0 +1,80 @@
+from typing import Any, Dict, Tuple
+import torchvision.transforms as transforms
+
+
+class PtvTransformsDataAugment:
+    """
+    Applies data augmentation transformations to dataset elements.
+    Supported transformations include flipping, rotation, translation, and scaling.
+
+    category: PyTorch wrapper - Training
+    """
+
+    @classmethod
+    def INPUT_TYPES(cls) -> Dict[str, Any]:
+        """
+        Defines the required input types for the transformations.
+
+        Returns:
+            Dict[str, Any]: A dictionary specifying the required input types and their constraints.
+        """
+        return {
+            "required": {
+                "h_flip_prob": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.001}),
+                "v_flip_prob": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.001}),
+                "rotate_degree": ("FLOAT", {"default": 0.0, "min": -360.0, "max": 360.0, "step": 0.001}),
+                "h_translate_ratio": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.001}),
+                "v_translate_ratio": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.001}),
+                "min_scale": ("FLOAT", {"default": 1.0, "min": 0.5, "max": 1.0, "step": 0.001}),
+                "max_scale": ("FLOAT", {"default": 1.0, "min": 1.0, "max": 2.0, "step": 0.001})
+            },
+        }
+
+    RETURN_TYPES: Tuple[str] = ("PTVTRANSFORM",)
+    FUNCTION: str = "f"
+    CATEGORY: str = "Training"
+
+    def f(self,
+          h_flip_prob: float,
+          v_flip_prob: float,
+          rotate_degree: float,
+          h_translate_ratio: float,
+          v_translate_ratio: float,
+          min_scale: float,
+          max_scale: float) -> Tuple[transforms.Compose]:
+        """
+        Applies the specified transformations to the dataset.
+
+        Args:
+            h_flip_prob (float): Probability of applying horizontal flip.
+            v_flip_prob (float): Probability of applying vertical flip.
+            rotate_degree (float): Maximum rotation degree.
+            h_translate_ratio (float): Horizontal translation ratio.
+            v_translate_ratio (float): Vertical translation ratio.
+            min_scale (float): Minimum scale factor.
+            max_scale (float): Maximum scale factor.
+
+        Returns:
+            Tuple[transforms.Compose]: A tuple containing the transformation pipeline.
+        """
+        translate_tuple = (h_translate_ratio, v_translate_ratio) if h_translate_ratio > 0.0 or v_translate_ratio > 0.0 else None
+        scale_tuple = (min_scale, max_scale) if min_scale != 1.0 or max_scale != 1.0 else None
+
+        actions = []
+        if h_flip_prob > 0.0:
+            actions.append(transforms.RandomHorizontalFlip(p=h_flip_prob))
+        if v_flip_prob > 0.0:
+            actions.append(transforms.RandomVerticalFlip(p=v_flip_prob))
+
+        actions.append(
+            transforms.RandomAffine(
+                degrees=rotate_degree,
+                translate=translate_tuple,
+                scale=scale_tuple
+            )
+        )
+        actions.append(transforms.ToTensor())  # Converts image to tensor and scales to [0, 1]
+
+        transform_pipeline = transforms.Compose(actions)
+
+        return (transform_pipeline,)

pyproject.toml

@@ -1,7 +1,7 @@
 [project]
 name = "pt-wrapper"
 description = "PyTorch extension for ComfyUI featuring extensive PyTorch wrapper nodes for seamless tensor operations and PyTorch model training."
-version = "3.0.4"
+version = "3.0.5"
 license = {file = "LICENSE"}
 
 [project.urls]

tests/pytorch_wrapper/ptv_transforms_data_augment_test.py

@@ -0,0 +1,51 @@
+import os
+import sys
+import unittest
+import torch
+from PIL import Image
+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.ptv_transforms_data_augment import PtvTransformsDataAugment
+
+
+class TestPtvTransformsResize(unittest.TestCase):
+    def setUp(self):
+        """Set up test data."""
+        self.compose_node = PtvTransformsDataAugment()
+        self.h_flip_prob = 0.4
+        self.v_flip_prob = 0.6
+        self.rotate_degree = 10
+        self.h_translate_ratio = 0.1
+        self.v_translate_ratio = 0.15
+        self.min_scale = 0.9
+        self.max_scale = 1.1
+
+    def test_with_transform(self):
+        """Tests the transformation output."""
+        transform = self.compose_node.f(
+            self.h_flip_prob,
+            self.v_flip_prob,
+            self.rotate_degree,
+            self.h_translate_ratio,
+            self.v_translate_ratio,
+            self.min_scale,
+            self.max_scale
+        )[0]
+        self.assertTrue(callable(transform), "Transform should be callable.")
+
+        # Create a dummy image
+        img = Image.fromarray(np.random.randint(0, 256, (256, 256, 3), dtype=np.uint8))
+        transformed_img = transform(img)
+
+        # Check output type
+        self.assertIsInstance(transformed_img, torch.Tensor, "Output should be a PyTorch tensor.")
+
+        # Check output shape
+        self.assertEqual(transformed_img.shape, (3, img.height, img.width), "Output shape should match target dimensions.")
+
+
+if __name__ == "__main__":
+    unittest.main()