diff --git a/README.md b/README.md index deae203..c9c0469 100644 --- a/README.md +++ b/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. diff --git a/__init__.py b/__init__.py index c38a468..a32bf64 100644 --- a/__init__.py +++ b/__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", } diff --git a/docs/images/resnet_eval.png b/docs/images/resnet_eval.png new file mode 100644 index 0000000..7eb4072 Binary files /dev/null and b/docs/images/resnet_eval.png differ diff --git a/docs/images/resnet_train.png b/docs/images/resnet_train.png new file mode 100644 index 0000000..e9955f5 Binary files /dev/null and b/docs/images/resnet_train.png differ diff --git a/docs/model_training.md b/docs/model_training.md index d79a588..9d65c19 100644 --- a/docs/model_training.md +++ b/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) diff --git a/docs/reference/node_reference.md b/docs/reference/node_reference.md index 0027797..407faac 100644 --- a/docs/reference/node_reference.md +++ b/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 diff --git a/docs/reference/ptv_transforms_data_augment.md b/docs/reference/ptv_transforms_data_augment.md new file mode 100644 index 0000000..12dc527 --- /dev/null +++ b/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 | + +
+Category: PyTorch wrapper - Training + +ComfyUI Data Analysis Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved. diff --git a/examples/workflows/cifar10_eval_v8.json b/examples/workflows/cifar10_eval_v8.json new file mode 100644 index 0000000..ee4030b --- /dev/null +++ b/examples/workflows/cifar10_eval_v8.json @@ -0,0 +1 @@ +{"last_node_id":89,"last_link_id":103,"nodes":[{"id":54,"type":"MarkdownNote","pos":[1712.6656494140625,2128.33642578125],"size":[210,58.94380569458008],"flags":{},"order":0,"mode":0,"inputs":[],"outputs":[],"properties":{},"widgets_values":["# Load val data"],"color":"#432","bgcolor":"#653"},{"id":78,"type":"PtvDatasetLen","pos":[2045.1402587890625,2289.77490234375],"size":[210,26],"flags":{},"order":8,"mode":0,"inputs":[{"name":"dataset","type":"PTVDATASET","link":86}],"outputs":[{"name":"INT","type":"INT","links":[85],"slot_index":0}],"properties":{"Node name for S&R":"PtvDatasetLen"},"widgets_values":[]},{"id":79,"type":"CDAShowInt","pos":[2045.1402587890625,2366.537353515625],"size":[315,76],"flags":{},"order":10,"mode":0,"inputs":[{"name":"int_scalar","type":"INT","link":85,"widget":{"name":"int_scalar"}}],"outputs":[{"name":"STRING","type":"STRING","links":null,"shape":6}],"properties":{"Node name for S&R":"CDAShowInt"},"widgets_values":[0,"10000"]},{"id":62,"type":"PtvDataset","pos":[2045.1402587890625,2050.991943359375],"size":[400,200],"flags":{},"order":5,"mode":0,"inputs":[{"name":"transform","type":"PTVTRANSFORM","link":89,"shape":7}],"outputs":[{"name":"PTVDATASET","type":"PTVDATASET","links":[75,86],"slot_index":0}],"properties":{"Node name for S&R":"PtvDataset"},"widgets_values":["CIFAR10",true,"","{\"train\":False}"]},{"id":84,"type":"CDAShowFloat","pos":[3364.6611328125,2185.21142578125],"size":[210,76],"flags":{},"order":13,"mode":0,"inputs":[{"name":"float_scalar","type":"FLOAT","link":95,"widget":{"name":"float_scalar"}}],"outputs":[{"name":"STRING","type":"STRING","links":null,"shape":6}],"title":"Recall","properties":{"Node name for S&R":"CDAShowFloat"},"widgets_values":[0,"0.9311999999999999"]},{"id":85,"type":"CDAShowFloat","pos":[3637.660888671875,2183.1708984375],"size":[210,76],"flags":{},"order":14,"mode":0,"inputs":[{"name":"float_scalar","type":"FLOAT","link":96,"widget":{"name":"float_scalar"}}],"outputs":[{"name":"STRING","type":"STRING","links":null,"shape":6}],"title":"F1","properties":{"Node name for S&R":"CDAShowFloat"},"widgets_values":[0,"0.9313751833176991"]},{"id":83,"type":"CDAShowInt","pos":[3364.6611328125,2323.3916015625],"size":[210,79.09933471679688],"flags":{},"order":15,"mode":0,"inputs":[{"name":"int_scalar","type":"INT","link":94,"widget":{"name":"int_scalar"}}],"outputs":[{"name":"STRING","type":"STRING","links":null,"shape":6}],"title":"Num Samples","properties":{"Node name for S&R":"CDAShowInt"},"widgets_values":[0,"10000"]},{"id":81,"type":"PtvTransformsToTensor","pos":[1708.258544921875,2050.991943359375],"size":[302.4000244140625,26],"flags":{},"order":1,"mode":0,"inputs":[],"outputs":[{"name":"PTVTRANSFORM","type":"PTVTRANSFORM","links":[89],"slot_index":0}],"properties":{"Node name for S&R":"PtvTransformsToTensor"},"widgets_values":[]},{"id":87,"type":"CDAShowFloat","pos":[3364.6611328125,2050.991943359375],"size":[210,76],"flags":{},"order":11,"mode":0,"inputs":[{"name":"float_scalar","type":"FLOAT","link":98,"widget":{"name":"float_scalar"}}],"outputs":[{"name":"STRING","type":"STRING","links":null,"shape":6}],"title":"Accuracy","properties":{"Node name for S&R":"CDAShowFloat"},"widgets_values":[0,"0.9312"]},{"id":86,"type":"CDAShowFloat","pos":[3637.660888671875,2050.991943359375],"size":[210,76],"flags":{},"order":12,"mode":0,"inputs":[{"name":"float_scalar","type":"FLOAT","link":97,"widget":{"name":"float_scalar"}}],"outputs":[{"name":"STRING","type":"STRING","links":null,"shape":6}],"title":"Precision","properties":{"Node name for S&R":"CDAShowFloat"},"widgets_values":[0,"0.931929003822688"]},{"id":88,"type":"PtLoadModel","pos":[2973.5009765625,1907.97412109375],"size":[303.63604736328125,75.06120300292969],"flags":{},"order":6,"mode":0,"inputs":[{"name":"model","type":"PTMODEL","link":101},{"name":"model_path","type":"STRING","link":100,"widget":{"name":"model_path"}}],"outputs":[{"name":"model","type":"PTMODEL","links":[93],"slot_index":0}],"properties":{"Node name for S&R":"PtLoadModel"},"widgets_values":["/container/transfer/models/food101_conv_ep5_v2.pt"]},{"id":82,"type":"PtEvaluateClassificationModel","pos":[2968.29638671875,2050.991943359375],"size":[313.0364990234375,146.49288940429688],"flags":{},"order":9,"mode":0,"inputs":[{"name":"model","type":"PTMODEL","link":93},{"name":"data_loader","type":"PTDATALOADER","link":102}],"outputs":[{"name":"accuracy","type":"FLOAT","links":[98],"slot_index":0},{"name":"precision","type":"FLOAT","links":[97],"slot_index":1},{"name":"recall","type":"FLOAT","links":[95],"slot_index":2},{"name":"f1","type":"FLOAT","links":[96],"slot_index":3},{"name":"num_samples","type":"INT","links":[94],"slot_index":4}],"properties":{"Node name for S&R":"PtEvaluateClassificationModel"},"widgets_values":[true]},{"id":80,"type":"PtnResnetModel","pos":[2498.54833984375,1782.458740234375],"size":[315,106],"flags":{},"order":2,"mode":0,"inputs":[],"outputs":[{"name":"PTMODEL","type":"PTMODEL","links":[101],"slot_index":0}],"properties":{"Node name for S&R":"PtnResnetModel"},"widgets_values":["(3,32,32)",10,10]},{"id":52,"type":"MarkdownNote","pos":[2498.54833984375,1657.6710205078125],"size":[210,58.94380569458008],"flags":{},"order":3,"mode":0,"inputs":[],"outputs":[],"properties":{},"widgets_values":["# Define model"],"color":"#432","bgcolor":"#653"},{"id":89,"type":"CDAStringCreate","pos":[2968.29638671875,1782.458740234375],"size":[315,58],"flags":{},"order":4,"mode":0,"inputs":[],"outputs":[{"name":"STRING","type":"STRING","links":[100],"slot_index":0}],"properties":{"Node name for S&R":"CDAStringCreate"},"widgets_values":["/container/transfer/models/cifar10_v8.pt"]},{"id":64,"type":"PtDataLoader","pos":[2498.54833984375,2050.991943359375],"size":[313.175048828125,152.2939453125],"flags":{},"order":7,"mode":0,"inputs":[{"name":"dataset","type":"PTVDATASET","link":75}],"outputs":[{"name":"PTDATALOADER","type":"PTDATALOADER","links":[102],"slot_index":0}],"properties":{"Node name for S&R":"PtDataLoader"},"widgets_values":[64,false,"{\"num_workers\":1}"]}],"links":[[75,62,0,64,0,"PTVDATASET"],[85,78,0,79,0,"INT"],[86,62,0,78,0,"PTVDATASET"],[89,81,0,62,0,"PTVTRANSFORM"],[93,88,0,82,0,"PTMODEL"],[94,82,4,83,0,"INT"],[95,82,2,84,0,"FLOAT"],[96,82,3,85,0,"FLOAT"],[97,82,1,86,0,"FLOAT"],[98,82,0,87,0,"FLOAT"],[100,89,0,88,1,"STRING"],[101,80,0,88,0,"PTMODEL"],[102,64,0,82,1,"PTDATALOADER"]],"groups":[],"config":{},"extra":{"ds":{"scale":0.794004118339972,"offset":[-2365.45583450116,-1763.5543973609274]},"node_versions":{"pt-wrapper":"f01f26eaa71112302f3512b2a5e31329f29fdd84","data-analysis":"bde4a07755390554f858e0b8fd8554a3123a6612"},"ue_links":[]},"version":0.4} \ No newline at end of file diff --git a/examples/workflows/cifar10_train_v8.json b/examples/workflows/cifar10_train_v8.json new file mode 100644 index 0000000..b9652ea --- /dev/null +++ b/examples/workflows/cifar10_train_v8.json @@ -0,0 +1 @@ +{"last_node_id":90,"last_link_id":103,"nodes":[{"id":61,"type":"PtvDataset","pos":[2007.994873046875,1723.90625],"size":[400,200],"flags":{},"order":9,"mode":0,"inputs":[{"name":"transform","type":"PTVTRANSFORM","link":103,"shape":7}],"outputs":[{"name":"PTVDATASET","type":"PTVDATASET","links":[74,83],"slot_index":0}],"properties":{"Node name for S&R":"PtvDataset"},"widgets_values":["CIFAR10",true,"","{\"train\":True}"]},{"id":19,"type":"PtSaveModel","pos":[3530.668701171875,1647.40185546875],"size":[396.8008117675781,75.65961456298828],"flags":{},"order":19,"mode":0,"inputs":[{"name":"model","type":"PTMODEL","link":90},{"name":"model_path","type":"STRING","link":99,"widget":{"name":"model_path"}}],"outputs":[{"name":"model full path","type":"STRING","links":[70],"slot_index":0}],"properties":{"Node name for S&R":"PtSaveModel"},"widgets_values":["/container/transfer/models/cifar10.pt"]},{"id":50,"type":"PtTrainClassificationModel","pos":[3138.080078125,1721.2340087890625],"size":[243.70773315429688,214],"flags":{},"order":17,"mode":0,"inputs":[{"name":"model","type":"PTMODEL","link":92},{"name":"train_loader","type":"PTDATALOADER","link":76},{"name":"optimizer","type":"PTOPTIMIZER","link":68},{"name":"val_loader","type":"PTDATALOADER","link":77,"shape":7}],"outputs":[{"name":"Model","type":"PTMODEL","links":[90],"slot_index":0},{"name":"train loss","type":"TENSOR","links":null},{"name":"val loss","type":"TENSOR","links":null}],"properties":{"Node name for S&R":"PtTrainClassificationModel"},"widgets_values":[100,true,false,20,false]},{"id":62,"type":"PtvDataset","pos":[2007.994873046875,2039.8226318359375],"size":[400,200],"flags":{},"order":10,"mode":0,"inputs":[{"name":"transform","type":"PTVTRANSFORM","link":89,"shape":7}],"outputs":[{"name":"PTVDATASET","type":"PTVDATASET","links":[75,86],"slot_index":0}],"properties":{"Node name for S&R":"PtvDataset"},"widgets_values":["CIFAR10",true,"","{\"train\":False}"]},{"id":76,"type":"PtvDatasetLen","pos":[2007.994873046875,1638.4072265625],"size":[210,26],"flags":{},"order":13,"mode":0,"inputs":[{"name":"dataset","type":"PTVDATASET","link":83}],"outputs":[{"name":"INT","type":"INT","links":[84],"slot_index":0}],"properties":{"Node name for S&R":"PtvDatasetLen"},"widgets_values":[]},{"id":78,"type":"PtvDatasetLen","pos":[2007.994873046875,2297.724365234375],"size":[210,26],"flags":{},"order":15,"mode":0,"inputs":[{"name":"dataset","type":"PTVDATASET","link":86}],"outputs":[{"name":"INT","type":"INT","links":[85],"slot_index":0}],"properties":{"Node name for S&R":"PtvDatasetLen"},"widgets_values":[]},{"id":79,"type":"CDAShowInt","pos":[2007.994873046875,2381.43359375],"size":[315,76],"flags":{},"order":18,"mode":0,"inputs":[{"name":"int_scalar","type":"INT","link":85,"widget":{"name":"int_scalar"}}],"outputs":[{"name":"STRING","type":"STRING","links":null,"shape":6}],"properties":{"Node name for S&R":"CDAShowInt"},"widgets_values":[0,"10000"]},{"id":64,"type":"PtDataLoader","pos":[2473.048828125,2039.8226318359375],"size":[400,200],"flags":{},"order":14,"mode":0,"inputs":[{"name":"dataset","type":"PTVDATASET","link":75}],"outputs":[{"name":"PTDATALOADER","type":"PTDATALOADER","links":[77],"slot_index":0}],"properties":{"Node name for S&R":"PtDataLoader"},"widgets_values":[128,false,"{\"num_workers\":1}"]},{"id":53,"type":"MarkdownNote","pos":[1215.7994384765625,1723.90625],"size":[210,58.94380569458008],"flags":{},"order":0,"mode":0,"inputs":[],"outputs":[],"properties":{},"widgets_values":["# Load train data"],"color":"#432","bgcolor":"#653"},{"id":54,"type":"MarkdownNote","pos":[1214.4312744140625,2039.8226318359375],"size":[210,58.94380569458008],"flags":{},"order":1,"mode":0,"inputs":[],"outputs":[],"properties":{},"widgets_values":["# Load val data"],"color":"#432","bgcolor":"#653"},{"id":77,"type":"CDAShowInt","pos":[2007.994873046875,1504.9381103515625],"size":[315,76],"flags":{},"order":16,"mode":0,"inputs":[{"name":"int_scalar","type":"INT","link":84,"widget":{"name":"int_scalar"}}],"outputs":[{"name":"STRING","type":"STRING","links":null,"shape":6}],"properties":{"Node name for S&R":"CDAShowInt"},"widgets_values":[0,"50000"]},{"id":90,"type":"PtvTransformsDataAugment","pos":[1595.48486328125,1723.90625],"size":[340.20001220703125,202],"flags":{},"order":2,"mode":0,"inputs":[],"outputs":[{"name":"PTVTRANSFORM","type":"PTVTRANSFORM","links":[103],"slot_index":0}],"properties":{"Node name for S&R":"PtvTransformsDataAugment"},"widgets_values":[0.5,0,0,0.125,0.125,1,1]},{"id":81,"type":"PtvTransformsToTensor","pos":[1595.48486328125,2039.8226318359375],"size":[338.8136291503906,26],"flags":{},"order":3,"mode":0,"inputs":[],"outputs":[{"name":"PTVTRANSFORM","type":"PTVTRANSFORM","links":[89],"slot_index":0}],"properties":{"Node name for S&R":"PtvTransformsToTensor"},"widgets_values":[]},{"id":63,"type":"PtDataLoader","pos":[2473.048828125,1723.90625],"size":[400,200],"flags":{},"order":12,"mode":0,"inputs":[{"name":"dataset","type":"PTVDATASET","link":74}],"outputs":[{"name":"PTDATALOADER","type":"PTDATALOADER","links":[76],"slot_index":0}],"properties":{"Node name for S&R":"PtDataLoader"},"widgets_values":[128,true,"{\"num_workers\":1}"]},{"id":55,"type":"MarkdownNote","pos":[3139.32958984375,1987.8914794921875],"size":[232.82992553710938,58],"flags":{},"order":4,"mode":0,"inputs":[],"outputs":[],"properties":{},"widgets_values":["# Train"],"color":"#432","bgcolor":"#653"},{"id":56,"type":"MarkdownNote","pos":[3535.058349609375,1775.04931640625],"size":[210,58.94380569458008],"flags":{},"order":5,"mode":0,"inputs":[],"outputs":[],"properties":{},"widgets_values":["# Save model"],"color":"#432","bgcolor":"#653"},{"id":51,"type":"PandasShowText","pos":[3536.080078125,1878.28125],"size":[333.6095275878906,79.25810241699219],"flags":{},"order":20,"mode":0,"inputs":[{"name":"text","type":"STRING","link":70,"widget":{"name":"text"}}],"outputs":[{"name":"STRING","type":"STRING","links":null,"shape":6}],"properties":{"Node name for S&R":"PandasShowText"},"widgets_values":["","/container/transfer/models/cifar10_v8.pt"]},{"id":52,"type":"MarkdownNote","pos":[2485.625732421875,1294.277587890625],"size":[210,58.94380569458008],"flags":{},"order":6,"mode":0,"inputs":[],"outputs":[],"properties":{},"widgets_values":["# Define model"],"color":"#432","bgcolor":"#653"},{"id":11,"type":"PtoAdam","pos":[3133.25634765625,1413.751953125],"size":[246.16842651367188,110.40398406982422],"flags":{},"order":11,"mode":0,"inputs":[{"name":"model","type":"PTMODEL","link":91}],"outputs":[{"name":"PTOPTIMIZER","type":"PTOPTIMIZER","links":[68],"slot_index":0}],"properties":{"Node name for S&R":"PtoAdam"},"widgets_values":[0.001,0.9,0.999]},{"id":80,"type":"PtnResnetModel","pos":[2478.35009765625,1413.751953125],"size":[379.3387145996094,113.2640609741211],"flags":{},"order":7,"mode":0,"inputs":[],"outputs":[{"name":"PTMODEL","type":"PTMODEL","links":[91,92],"slot_index":0}],"properties":{"Node name for S&R":"PtnResnetModel"},"widgets_values":["(3,32,32)",10,10]},{"id":89,"type":"CDAStringCreate","pos":[3536.05712890625,1534.6680908203125],"size":[385.5651550292969,59.037696838378906],"flags":{},"order":8,"mode":0,"inputs":[],"outputs":[{"name":"STRING","type":"STRING","links":[99],"slot_index":0}],"properties":{"Node name for S&R":"CDAStringCreate"},"widgets_values":["/container/transfer/models/cifar10_v8.pt"]}],"links":[[68,11,0,50,2,"PTOPTIMIZER"],[70,19,0,51,0,"STRING"],[74,61,0,63,0,"PTVDATASET"],[75,62,0,64,0,"PTVDATASET"],[76,63,0,50,1,"PTDATALOADER"],[77,64,0,50,3,"PTDATALOADER"],[83,61,0,76,0,"PTVDATASET"],[84,76,0,77,0,"INT"],[85,78,0,79,0,"INT"],[86,62,0,78,0,"PTVDATASET"],[89,81,0,62,0,"PTVTRANSFORM"],[90,50,0,19,0,"PTMODEL"],[91,80,0,11,0,"PTMODEL"],[92,80,0,50,0,"PTMODEL"],[99,89,0,19,1,"STRING"],[103,90,0,61,0,"PTVTRANSFORM"]],"groups":[],"config":{},"extra":{"ds":{"scale":0.41208717708376225,"offset":[-514.8735211991554,-1049.4380234927548]},"node_versions":{"pt-wrapper":"f01f26eaa71112302f3512b2a5e31329f29fdd84","data-analysis":"bde4a07755390554f858e0b8fd8554a3123a6612"},"ue_links":[]},"version":0.4} \ No newline at end of file diff --git a/modules/pytorch_wrapper/ptv_transforms_data_augment.py b/modules/pytorch_wrapper/ptv_transforms_data_augment.py new file mode 100644 index 0000000..f549e33 --- /dev/null +++ b/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,) diff --git a/pyproject.toml b/pyproject.toml index 3d6e1eb..3796f4b 100644 --- a/pyproject.toml +++ b/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] diff --git a/tests/pytorch_wrapper/ptv_transforms_data_augment_test.py b/tests/pytorch_wrapper/ptv_transforms_data_augment_test.py new file mode 100644 index 0000000..2f97adc --- /dev/null +++ b/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()