diff --git a/README.md b/README.md index 6fb626b..113a323 100644 --- a/README.md +++ b/README.md @@ -1,45 +1,10 @@ # ComfyUI-Data-Analysis -## Update (March 5, 2025) -Training workflow trains the FashionMNIST model from scratch with a single click of a button. +## Update (March 7, 2025) +You can now use newly added nodes to train a machine learning model for image classification with your own image data **without any coding** from scratch. -Experimental PyTorch model training has been added (documentation coming soon). -No coding required—training is as simple as a single click! - -It is designed to be modular. For example, to switch from a linear model to a conv model, you can just replace the model node (see screenshot below). - -Or you can just drag and drop below workflows to ComfyUI to check it out. Model is downloaded automatically: - -Linear model -* Training: examples/fashion_mnist_train.json -* Evaluation: examples/fashion_mnist_eval.json - -Convolutional model -* Training: examples/fashion_mnist_train_conv.json -* Evaluation: examples/fashion_mnist_eval_conv.json - -These training workflows train a FashionMNIST model from scratch effortlessly. - -### Training workflow -**Linear model** -![Train](docs/images/fashion_mnist_train.png) - -**Conv model** -![Train](docs/images/conv_train.png) - -*Close up of the conv node* -![Train](docs/images/conv_train2.png) - -### Eval workflow -**Linear model** -![Eval](docs/images/fashion_mnist_eval.png) - -**Conv model** -![Eval](docs/images/conv_eval.png) - -Note that the training feature is a proof of concept and currently supports only the linear and conv models. Additional models may be added based on demand. - -
+Check out [the new tutorial to classify dogs and cats](docs/dog_cat_classification_model_training.md) +--- ![Baseball tutorial](docs/images/baseball_question.jpg) ## How Can We Figure Out Which MLB Player Had the Most Hits Per Year? diff --git a/__init__.py b/__init__.py index 4ea7d6b..6276f3c 100644 --- a/__init__.py +++ b/__init__.py @@ -243,6 +243,7 @@ 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_resize import PtvTransformsResize 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 @@ -494,6 +495,7 @@ NODE_CLASS_MAPPINGS: Dict[str, Type[T]] = { "PtvDataset": PtvDataset, "PtvDatasetLoader": PtvDatasetLoader, "PtvImageFolderDataset": PtvImageFolderDataset, + "PtvTransformsResize": PtvTransformsResize, "PtvTransformsToTensor": PtvTransformsToTensor, "PyDatetimeToString": PyDatetimeToString, "PyFloatToString": PyFloatToString, @@ -754,6 +756,7 @@ NODE_DISPLAY_NAME_MAPPINGS: Dict[str, str] = { "PtvDataset": "Ptv Dataset", "PtvDatasetLoader": "Ptv Dataset Loader", "PtvImageFolderDataset": "Ptv Image Folder Dataset", + "PtvTransformsResize": "Ptv Transforms Resize", "PtvTransformsToTensor": "Ptv Transforms To Tensor", "PyDatetimeToString": "Py Datetime To String", "PyFloatToString": "Py Float To String", diff --git a/dev/scripts/run_tests.sh b/dev/scripts/run_tests.sh index cd29541..9c295c8 100755 --- a/dev/scripts/run_tests.sh +++ b/dev/scripts/run_tests.sh @@ -8,5 +8,5 @@ # the tests in the directory will be skipped. export PYTHONPATH=modules:$PYTHONPATH echo $PYTHONPATH -python -m unittest discover -s tests -p "*_test.py" -t . +RUN_SKIPPED_TESTS=0 python -m unittest discover -s tests -p "*_test.py" -t . diff --git a/docs/dog_cat_classification_model_training.md b/docs/dog_cat_classification_model_training.md new file mode 100644 index 0000000..68f3c13 --- /dev/null +++ b/docs/dog_cat_classification_model_training.md @@ -0,0 +1,137 @@ +# Dog cat classification model training tutorial +This tutorial covers how to use **ComfyUI Data Analysis** to train an image classification model with your own image dataset. +We'll use a **dog and cat image dataset** created specifically for this tutorial, which is available publicly on [my GitHub repo](https://github.com/HowToSD/dog_and_cat_dataset). + +Once you complete the tutorial, you can experiment with your own images. + +**Disclaimer:** The model architecture is a traditional convolutional neural network (CNN) and is **not state-of-the-art**. This choice keeps the model small, enabling faster training and easier explanation. + +--- + +## Overview of Steps + +1. Downloading the dog and cat dataset from the GitHub repo +2. Edit the training workflow +3. Run the training workflow +4. Edit the evaluation workflow +5. Run the evaluation workflow + +--- + +## 1. Download the Dog & Cat Dataset + +Visit [the GitHub repo](https://github.com/HowToSD/dog_and_cat_dataset) and go to the **Releases** page. +Click **"Source code (zip)"** to download the zip that contains all the images. +Extract the ZIP file using an unarchiver. + +The dataset follows this directory structure: + +``` +dog_and_cat + train + cats + dogs + + val + cats + dogs +``` + + +Take note of the **full paths** to the `train` and `val` directories, as you'll need them in the next step. + +--- + +## 2. Editing the Training Workflow + +Drag and drop the below file into ComfyUI: +``` +examples/workflow/dog_cat_classification_train_model.json +``` +Some nodes will be marked red, indicating they require edits. +The only required changes are the **train** and **val** dataset paths. +Replace them with the paths noted earlier. + +--- + +## 3. Running the Training Workflow + +Click **Queue** to start training. + +Monitor the console for errors. If training runs smoothly, you should see output like this: + +``` +Epoch (train) 1/20, Loss: 0.6865 +Epoch (val) 1/20, Loss: 0.6801 +Epoch (train) 2/20, Loss: 0.6216 +Epoch (val) 2/20, Loss: 0.5710 +Epoch (train) 3/20, Loss: 0.4998 +Epoch (val) 3/20, Loss: 0.4716 +Epoch (train) 4/20, Loss: 0.3827 +Epoch (val) 4/20, Loss: 0.4045 +Epoch (train) 5/20, Loss: 0.2897 +Epoch (val) 5/20, Loss: 0.3297 +Epoch (train) 6/20, Loss: 0.2243 +Epoch (val) 6/20, Loss: 0.2925 +Epoch (train) 7/20, Loss: 0.1603 +Epoch (val) 7/20, Loss: 0.2916 +Epoch (train) 8/20, Loss: 0.1102 +Epoch (val) 8/20, Loss: 0.3216 +Epoch (train) 9/20, Loss: 0.0916 +Epoch (val) 9/20, Loss: 0.2706 +Epoch (train) 10/20, Loss: 0.0522 +Epoch (val) 10/20, Loss: 0.2220 +Epoch (train) 11/20, Loss: 0.0481 +Epoch (val) 11/20, Loss: 0.2820 +Epoch (train) 12/20, Loss: 0.0228 +Epoch (val) 12/20, Loss: 0.2899 +Epoch (train) 13/20, Loss: 0.0130 +Epoch (val) 13/20, Loss: 0.3056 +Epoch (train) 14/20, Loss: 0.0137 +Epoch (val) 14/20, Loss: 0.3070 +Epoch (train) 15/20, Loss: 0.0200 +Epoch (val) 15/20, Loss: 0.2714 +Epoch (train) 16/20, Loss: 0.0313 +Epoch (val) 16/20, Loss: 0.3108 +Epoch (train) 17/20, Loss: 0.0125 +Epoch (val) 17/20, Loss: 0.3388 +Epoch (train) 18/20, Loss: 0.0174 +Epoch (val) 18/20, Loss: 0.2989 +Epoch (train) 19/20, Loss: 0.0017 +Epoch (val) 19/20, Loss: 0.3165 +Epoch (train) 20/20, Loss: 0.0009 +Epoch (val) 20/20, Loss: 0.3582 +Prompt executed in 253.64 seconds +``` + +If you encounter an error: +- Ensure dataset paths are correct. +- Reduce **batch size** if your GPU runs out of memory (OOM error). + +Once training completes, a **line chart** displaying training progress will appear at the bottom right of the canvas. You don’t need to interpret it initially. + +--- + +## 4. Editing the Evaluation Workflow +Now, let's verify if the model was trained correctly. In this step, we'll feed the validation data into the trained model and compute metrics to assess its performance. + +Drag and drop the below file into ComfyUI: +``` +examples/workflow/dog_cat_classification_eval_model.json +``` + +into **ComfyUI**. + +Update the **val dataset path**. + +--- + +## 5. Running the Evaluation Workflow + +Click **Queue**. The evaluation process runs much faster than training. + +Check the **Accuracy** node—it should display a value around **0.9**, confirming successful model training. + +Congratulations! You’ve successfully trained a machine learning model from scratch without writing a single line of code. + +Now, you can use your own images to train the model in the same way. \ No newline at end of file diff --git a/docs/model_training.md b/docs/model_training.md new file mode 100644 index 0000000..dbfc555 --- /dev/null +++ b/docs/model_training.md @@ -0,0 +1,49 @@ +# Model Training + +Experimental PyTorch model training has been added. +No coding is required. Training is as simple as a single click! +Currently, only classification is supported. + +The system is designed to be modular. For example, to switch from a linear model to a convolutional model, you only need to replace the model node (see screenshot below). + +You can choose to: +* Use your own images +* Use a public dataset + +## Using Your Own Images + +Check out [the dog and cat classification tutorial](dog_cat_classification_model_training.md). + +## Public Datasets + +Currently, only Fashion MNIST has 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. + +Linear model +* Training: examples/fashion_mnist_train.json +* Evaluation: examples/fashion_mnist_eval.json + +Convolutional model +* Training: examples/fashion_mnist_train_conv.json +* Evaluation: examples/fashion_mnist_eval_conv.json + +These training workflows train a FashionMNIST model from scratch effortlessly. + +### Training workflow +**Linear model** +![Train](images/fashion_mnist_train.png) + +**Conv model** +![Train](images/conv_train.png) + +*Close up of the conv node* +![Train](images/conv_train2.png) + +### Eval workflow +**Linear model** +![Eval](images/fashion_mnist_eval.png) + +**Conv model** +![Eval](images/conv_eval.png) diff --git a/docs/reference/node_reference.md b/docs/reference/node_reference.md index 1cf3d9f..f115712 100644 --- a/docs/reference/node_reference.md +++ b/docs/reference/node_reference.md @@ -276,6 +276,7 @@ | [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 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 | Node | Description | diff --git a/docs/reference/ptv_transforms_resize.md b/docs/reference/ptv_transforms_resize.md new file mode 100644 index 0000000..159626f --- /dev/null +++ b/docs/reference/ptv_transforms_resize.md @@ -0,0 +1,18 @@ +# Ptv Transforms Resize +Resizes and transforms elements of dataset to PyTorch tensors. + +## Input +| Name | Data type | +|---|---| +| height | Int | +| width | Int | + +## 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/dog_cat_classification_eval_model.json b/examples/workflows/dog_cat_classification_eval_model.json new file mode 100644 index 0000000..aff7eff --- /dev/null +++ b/examples/workflows/dog_cat_classification_eval_model.json @@ -0,0 +1 @@ +{"last_node_id":54,"last_link_id":67,"nodes":[{"id":49,"type":"PtDataLoader","pos":[2664.0126953125,2004.834716796875],"size":[308.1214599609375,135.3099822998047],"flags":{},"order":5,"mode":0,"inputs":[{"name":"dataset","type":"PTVDATASET","link":64}],"outputs":[{"name":"PTDATALOADER","type":"PTDATALOADER","links":[65],"slot_index":0}],"properties":{"Node name for S&R":"PtDataLoader"},"widgets_values":[8,false,"{\"num_workers\":1}"]},{"id":50,"type":"PtvImageFolderDataset","pos":[2324.53857421875,2006.695068359375],"size":[315,58],"flags":{},"order":3,"mode":0,"inputs":[{"name":"transform","type":"PTVTRANSFORM","link":67,"shape":7}],"outputs":[{"name":"PTVDATASET","type":"PTVDATASET","links":[64],"slot_index":0}],"properties":{"Node name for S&R":"PtvImageFolderDataset"},"widgets_values":["/container/transfer/datasets/dog_and_cat/dog_and_cat/val/"]},{"id":53,"type":"PtvTransformsResize","pos":[1959.0155029296875,2007.09521484375],"size":[315,82],"flags":{},"order":0,"mode":0,"inputs":[],"outputs":[{"name":"PTVTRANSFORM","type":"PTVTRANSFORM","links":[67],"slot_index":0}],"properties":{"Node name for S&R":"PtvTransformsResize"},"widgets_values":[256,256]},{"id":48,"type":"PtnConvModel","pos":[2666.497314453125,1718.6485595703125],"size":[315,202],"flags":{},"order":1,"mode":0,"inputs":[],"outputs":[{"name":"PTMODEL","type":"PTMODEL","links":[62],"slot_index":0}],"properties":{"Node name for S&R":"PtnConvModel"},"widgets_values":["(3,256,256)",0,2,"[32,64,128,256,512]","[3,3,3,3,3]","[1,1,1,1,1]","[True,True,True,True,True]"]},{"id":47,"type":"PtEvaluateClassificationModel","pos":[3073.28466796875,1998.2191162109375],"size":[268.79998779296875,148.22767639160156],"flags":{},"order":6,"mode":0,"inputs":[{"name":"model","type":"PTMODEL","link":55},{"name":"data_loader","type":"PTDATALOADER","link":65}],"outputs":[{"name":"accuracy","type":"FLOAT","links":[66],"slot_index":0},{"name":"precision","type":"FLOAT","links":[58],"slot_index":1},{"name":"recall","type":"FLOAT","links":[59],"slot_index":2},{"name":"f1","type":"FLOAT","links":[60],"slot_index":3},{"name":"num_samples","type":"INT","links":[61],"slot_index":4}],"properties":{"Node name for S&R":"PtEvaluateClassificationModel"},"widgets_values":[true]},{"id":33,"type":"CDAShowInt","pos":[3507.768310546875,2136.57861328125],"size":[210,79.09933471679688],"flags":{},"order":11,"mode":0,"inputs":[{"name":"int_scalar","type":"INT","link":61,"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,"300"]},{"id":30,"type":"CDAShowFloat","pos":[3507.768310546875,1986.105712890625],"size":[210,76],"flags":{},"order":9,"mode":0,"inputs":[{"name":"float_scalar","type":"FLOAT","link":59,"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.9"]},{"id":31,"type":"CDAShowFloat","pos":[3768.4755859375,1980.9920654296875],"size":[210,76],"flags":{},"order":10,"mode":0,"inputs":[{"name":"float_scalar","type":"FLOAT","link":60,"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.8997817469155049"]},{"id":29,"type":"CDAShowFloat","pos":[3768.4755859375,1834.9833984375],"size":[210,76],"flags":{},"order":8,"mode":0,"inputs":[{"name":"float_scalar","type":"FLOAT","link":58,"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.9035150645624104"]},{"id":52,"type":"CDAShowFloat","pos":[3507.768310546875,1834.7109375],"size":[210,76],"flags":{},"order":7,"mode":0,"inputs":[{"name":"float_scalar","type":"FLOAT","link":66,"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.9"]},{"id":23,"type":"PtLoadModel","pos":[3075.703369140625,1723.319580078125],"size":[261.13433837890625,74.19380950927734],"flags":{},"order":4,"mode":0,"inputs":[{"name":"model","type":"PTMODEL","link":62}],"outputs":[{"name":"model","type":"PTMODEL","links":[55],"slot_index":0}],"properties":{"Node name for S&R":"PtLoadModel"},"widgets_values":["dog_cat.pt"]},{"id":54,"type":"MarkdownNote","pos":[2328.715576171875,1839.947509765625],"size":[313.0150451660156,99.62860870361328],"flags":{},"order":2,"mode":0,"inputs":[],"outputs":[],"properties":{},"widgets_values":["# Load val data\nYou need set the root field to point to the **val** directory of the dataset you downloaded."],"color":"#322","bgcolor":"#533"}],"links":[[55,23,0,47,0,"PTMODEL"],[58,47,1,29,0,"FLOAT"],[59,47,2,30,0,"FLOAT"],[60,47,3,31,0,"FLOAT"],[61,47,4,33,0,"INT"],[62,48,0,23,0,"PTMODEL"],[64,50,0,49,0,"PTVDATASET"],[65,49,0,47,1,"PTDATALOADER"],[66,47,0,52,0,"FLOAT"],[67,53,0,50,0,"PTVTRANSFORM"]],"groups":[],"config":{},"extra":{"ds":{"scale":0.34289106510190753,"offset":[-1247.259007394664,-447.540790015054]},"node_versions":{"data-analysis":"db92cda5f5a1e611e15ca6f2531ff8585a5607a5"},"ue_links":[]},"version":0.4} \ No newline at end of file diff --git a/examples/workflows/dog_cat_classification_train_model.json b/examples/workflows/dog_cat_classification_train_model.json new file mode 100644 index 0000000..706b774 --- /dev/null +++ b/examples/workflows/dog_cat_classification_train_model.json @@ -0,0 +1 @@ +{"last_node_id":86,"last_link_id":95,"nodes":[{"id":11,"type":"PtoAdam","pos":[3106.311767578125,1425.00830078125],"size":[265.8851013183594,108.32855224609375],"flags":{},"order":16,"mode":0,"inputs":[{"name":"model","type":"PTMODEL","link":89}],"outputs":[{"name":"PTOPTIMIZER","type":"PTOPTIMIZER","links":[68],"slot_index":0}],"properties":{"Node name for S&R":"PtoAdam"},"widgets_values":[0.0001,0.9,0.999]},{"id":72,"type":"PtDataLoader","pos":[2583.5888671875,1626.235595703125],"size":[400,200],"flags":{},"order":20,"mode":0,"inputs":[{"name":"dataset","type":"PTVDATASET","link":86}],"outputs":[{"name":"PTDATALOADER","type":"PTDATALOADER","links":[87],"slot_index":0}],"properties":{"Node name for S&R":"PtDataLoader"},"widgets_values":[8,true,"{\"num_workers\":1}"]},{"id":19,"type":"PtSaveModel","pos":[3464.489501953125,1694.661376953125],"size":[280.30267333984375,68.31287384033203],"flags":{},"order":23,"mode":0,"inputs":[{"name":"model","type":"PTMODEL","link":71}],"outputs":[{"name":"model full path","type":"STRING","links":[70],"slot_index":0}],"properties":{"Node name for S&R":"PtSaveModel"},"widgets_values":["dog_cat.pt"]},{"id":74,"type":"PtvImageFolderDataset","pos":[2228.7578125,1923.7598876953125],"size":[315,58],"flags":{},"order":18,"mode":0,"inputs":[{"name":"transform","type":"PTVTRANSFORM","link":95,"shape":7}],"outputs":[{"name":"PTVDATASET","type":"PTVDATASET","links":[91],"slot_index":0}],"properties":{"Node name for S&R":"PtvImageFolderDataset"},"widgets_values":["/container/transfer/datasets/dog_and_cat/dog_and_cat/val/"]},{"id":71,"type":"PtnConvModel","pos":[2607.428955078125,1271.409912109375],"size":[315,202],"flags":{},"order":0,"mode":0,"inputs":[],"outputs":[{"name":"PTMODEL","type":"PTMODEL","links":[88,89],"slot_index":0}],"properties":{"Node name for S&R":"PtnConvModel"},"widgets_values":["(3,256,256)",0,2,"[32,64,128,256,512]","[3,3,3,3,3]","[1,1,1,1,1]","[True,True,True,True,True]"]},{"id":53,"type":"MarkdownNote","pos":[2219.2607421875,1485.9119873046875],"size":[310.71795654296875,94.28345489501953],"flags":{},"order":1,"mode":0,"inputs":[],"outputs":[],"properties":{},"widgets_values":["# Load train data\nYou need set the root field to point to the **train** directory of the dataset you downloaded."],"color":"#322","bgcolor":"#533"},{"id":76,"type":"MarkdownNote","pos":[2014.2506103515625,926.4641723632812],"size":[436.95391845703125,69.52714538574219],"flags":{},"order":2,"mode":0,"inputs":[],"outputs":[],"properties":{},"widgets_values":["# Dog Cat Classification Model Training"],"color":"#432","bgcolor":"#653"},{"id":78,"type":"PtvTransformsResize","pos":[1820.2008056640625,1781.384521484375],"size":[315,82],"flags":{},"order":3,"mode":0,"inputs":[],"outputs":[{"name":"PTVTRANSFORM","type":"PTVTRANSFORM","links":[94,95],"slot_index":0}],"properties":{"Node name for S&R":"PtvTransformsResize"},"widgets_values":[256,256]},{"id":52,"type":"MarkdownNote","pos":[2605.098876953125,1119.3177490234375],"size":[308.6709289550781,100.76033782958984],"flags":{},"order":4,"mode":0,"inputs":[],"outputs":[],"properties":{},"widgets_values":["# Model set up\n\nNo need to edit this node unless you are familiar with machine learning"],"color":"#232","bgcolor":"#353"},{"id":79,"type":"MarkdownNote","pos":[1816.8997802734375,1627.8802490234375],"size":[317.4327697753906,98.5698471069336],"flags":{},"order":5,"mode":0,"inputs":[],"outputs":[],"properties":{},"widgets_values":["# Dataset transformation\n\nNo need to edit this node unless you are familiar with machine learning. This resizes input data to 256px x 256px."],"color":"#232","bgcolor":"#353"},{"id":75,"type":"PtDataLoader","pos":[2588.29736328125,1922.6170654296875],"size":[400,200],"flags":{},"order":21,"mode":0,"inputs":[{"name":"dataset","type":"PTVDATASET","link":91}],"outputs":[{"name":"PTDATALOADER","type":"PTDATALOADER","links":[93],"slot_index":0}],"properties":{"Node name for S&R":"PtDataLoader"},"widgets_values":[8,false,"{\"num_workers\":1}"]},{"id":51,"type":"PandasShowText","pos":[3845.79296875,1686.58056640625],"size":[499.0397033691406,76],"flags":{},"order":26,"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":["","/home/dockeruser/comfy/ComfyUI/custom_nodes/data-analysis/models/dog_cat.pt"]},{"id":58,"type":"PandasCreateFromTensor","pos":[3487.738037109375,2425.4267578125],"size":[315,26],"flags":{},"order":25,"mode":0,"inputs":[{"name":"tens","type":"TENSOR","link":73}],"outputs":[{"name":"DATAFRAME","type":"DATAFRAME","links":[75],"slot_index":0}],"properties":{"Node name for S&R":"PandasCreateFromTensor"},"widgets_values":[]},{"id":64,"type":"PandasCreateFromTensor","pos":[3844.351806640625,2508.688720703125],"size":[315,26],"flags":{},"order":19,"mode":0,"inputs":[{"name":"tens","type":"TENSOR","link":78}],"outputs":[{"name":"DATAFRAME","type":"DATAFRAME","links":[77],"slot_index":0}],"properties":{"Node name for S&R":"PandasCreateFromTensor"},"widgets_values":[]},{"id":65,"type":"PandasHorizontalConcat","pos":[4218.2119140625,2380.0498046875],"size":[315,78],"flags":{},"order":28,"mode":0,"inputs":[{"name":"a_dataframe","type":"DATAFRAME","link":79},{"name":"b_dataframe","type":"DATAFRAME","link":77}],"outputs":[{"name":"DATAFRAME","type":"DATAFRAME","links":[80,81],"slot_index":0}],"properties":{"Node name for S&R":"PandasHorizontalConcat"},"widgets_values":[true]},{"id":61,"type":"PandasRenameAdvanced","pos":[4608.04443359375,2380.0498046875],"size":[238.48175048828125,131.16506958007812],"flags":{},"order":30,"mode":0,"inputs":[{"name":"dataframe","type":"DATAFRAME","link":81}],"outputs":[{"name":"DATAFRAME","type":"DATAFRAME","links":[82],"slot_index":0}],"properties":{"Node name for S&R":"PandasRenameAdvanced"},"widgets_values":["columns","{0:\"train\",1:\"val\",2:\"epoch\"}"]},{"id":60,"type":"SNSLine","pos":[4918.4609375,2380.0498046875],"size":[315,250],"flags":{},"order":31,"mode":0,"inputs":[{"name":"dataframe","type":"DATAFRAME","link":82}],"outputs":[{"name":"IMAGE","type":"IMAGE","links":[84],"slot_index":0}],"properties":{"Node name for S&R":"SNSLine"},"widgets_values":["epoch","train,val","Loss","Epoch","Cross entropy loss","train,val",true,"",""]},{"id":68,"type":"SaveImage","pos":[5326.560546875,2380.0498046875],"size":[315,270],"flags":{},"order":32,"mode":0,"inputs":[{"name":"images","type":"IMAGE","link":84}],"outputs":[],"properties":{"Node name for S&R":"SaveImage"},"widgets_values":["dog_cat_train"]},{"id":57,"type":"PandasCreateFromTensor","pos":[3483.9892578125,2346.20166015625],"size":[315,26],"flags":{},"order":24,"mode":0,"inputs":[{"name":"tens","type":"TENSOR","link":72}],"outputs":[{"name":"DATAFRAME","type":"DATAFRAME","links":[74],"slot_index":0}],"properties":{"Node name for S&R":"PandasCreateFromTensor"},"widgets_values":[]},{"id":83,"type":"MarkdownNote","pos":[2943.528564453125,925.8107299804688],"size":[336.614990234375,70.34814453125],"flags":{},"order":6,"mode":0,"inputs":[],"outputs":[],"properties":{},"widgets_values":["## Yellow nodes may require adjustment"],"color":"#432","bgcolor":"#653"},{"id":80,"type":"MarkdownNote","pos":[3106.311767578125,1250.9371337890625],"size":[257.6202392578125,106.56401062011719],"flags":{},"order":7,"mode":2,"inputs":[],"outputs":[],"properties":{},"widgets_values":["# Optimizer\nYou may need to adjust the learning rate but starts with the number below."],"color":"#432","bgcolor":"#653"},{"id":50,"type":"PtTrainClassificationModel","pos":[3106.311767578125,1663.041748046875],"size":[265.619873046875,145.1302490234375],"flags":{},"order":22,"mode":0,"inputs":[{"name":"model","type":"PTMODEL","link":88},{"name":"train_loader","type":"PTDATALOADER","link":87},{"name":"optimizer","type":"PTOPTIMIZER","link":68},{"name":"val_loader","type":"PTDATALOADER","link":93,"shape":7}],"outputs":[{"name":"Model","type":"PTMODEL","links":[71],"slot_index":0},{"name":"train loss","type":"TENSOR","links":[72],"slot_index":1},{"name":"val loss","type":"TENSOR","links":[73],"slot_index":2}],"properties":{"Node name for S&R":"PtTrainClassificationModel"},"widgets_values":[20,true]},{"id":84,"type":"MarkdownNote","pos":[2596.4189453125,2210.809326171875],"size":[387.0498352050781,100.52718353271484],"flags":{},"order":8,"mode":2,"inputs":[],"outputs":[],"properties":{},"widgets_values":["# Data loading\nIf GPU throws OOM, make the batch_size smaller."],"color":"#432","bgcolor":"#653"},{"id":82,"type":"MarkdownNote","pos":[2944.823486328125,805.5132446289062],"size":[326.2669677734375,67.76112365722656],"flags":{},"order":9,"mode":0,"inputs":[],"outputs":[],"properties":{},"widgets_values":["## Red nodes require editing"],"color":"#322","bgcolor":"#533"},{"id":77,"type":"MarkdownNote","pos":[2230.19091796875,2034.1849365234375],"size":[313.0150451660156,99.62860870361328],"flags":{},"order":10,"mode":0,"inputs":[],"outputs":[],"properties":{},"widgets_values":["# Load val data\nYou need set the root field to point to the **val** directory of the dataset you downloaded."],"color":"#322","bgcolor":"#533"},{"id":62,"type":"PtArange","pos":[3481.524169921875,2515.711669921875],"size":[315,130],"flags":{},"order":11,"mode":0,"inputs":[],"outputs":[{"name":"TENSOR","type":"TENSOR","links":[78],"slot_index":0}],"properties":{"Node name for S&R":"PtArange"},"widgets_values":["0","20","","float32"],"color":"#432","bgcolor":"#653"},{"id":55,"type":"MarkdownNote","pos":[3106.311767578125,1861.1396484375],"size":[265.81390380859375,103.27005767822266],"flags":{},"order":12,"mode":2,"inputs":[],"outputs":[],"properties":{},"widgets_values":["# Train\nYou may need to adjust epochs but start with 20."],"color":"#432","bgcolor":"#653"},{"id":85,"type":"MarkdownNote","pos":[3485.391357421875,2688.060791015625],"size":[304.5123596191406,100.34938049316406],"flags":{},"order":13,"mode":2,"inputs":[],"outputs":[],"properties":{},"widgets_values":["# Update \"end\" field\nIf you change the epochs, change the end field to match."],"color":"#432","bgcolor":"#653"},{"id":66,"type":"PandasShowDataFrame","pos":[4215.47265625,2533.793701171875],"size":[315.237060546875,414.686279296875],"flags":{},"order":29,"mode":0,"inputs":[{"name":"dataframe","type":"DATAFRAME","link":80}],"outputs":[{"name":"STRING","type":"STRING","links":null,"shape":6}],"properties":{"Node name for S&R":"PandasShowDataFrame"},"widgets_values":[" 0 1 2\n0 0.686489 0.680127 0.0\n1 0.621646 0.571015 1.0\n2 0.499781 0.471638 2.0\n3 0.382664 0.404523 3.0\n4 0.289659 0.329657 4.0\n5 0.224341 0.292451 5.0\n6 0.160335 0.291596 6.0\n7 0.110199 0.321644 7.0\n8 0.091624 0.270642 8.0\n9 0.052231 0.221975 9.0\n10 0.048100 0.281967 10.0\n11 0.022782 0.289861 11.0\n12 0.013045 0.305584 12.0\n13 0.013724 0.307031 13.0\n14 0.020037 0.271379 14.0\n15 0.031278 0.310788 15.0\n16 0.012516 0.338837 16.0\n17 0.017392 0.298860 17.0\n18 0.001674 0.316514 18.0\n19 0.000855 0.358205 19.0"]},{"id":59,"type":"PandasHorizontalConcat","pos":[3849.372802734375,2380.0498046875],"size":[315,78],"flags":{},"order":27,"mode":0,"inputs":[{"name":"a_dataframe","type":"DATAFRAME","link":74},{"name":"b_dataframe","type":"DATAFRAME","link":75}],"outputs":[{"name":"DATAFRAME","type":"DATAFRAME","links":[79],"slot_index":0}],"properties":{"Node name for S&R":"PandasHorizontalConcat"},"widgets_values":[true]},{"id":86,"type":"MarkdownNote","pos":[3453.1630859375,2150.53857421875],"size":[1160.256103515625,64.83271026611328],"flags":{},"order":14,"mode":0,"inputs":[],"outputs":[],"properties":{},"widgets_values":["# This area is to show how training went. Initially, it's OK if you don't understand the meaning."],"color":"#232","bgcolor":"#353"},{"id":69,"type":"PtvImageFolderDataset","pos":[2220.0927734375,1629.250732421875],"size":[315,58],"flags":{},"order":17,"mode":0,"inputs":[{"name":"transform","type":"PTVTRANSFORM","link":94,"shape":7}],"outputs":[{"name":"PTVDATASET","type":"PTVDATASET","links":[86],"slot_index":0}],"properties":{"Node name for S&R":"PtvImageFolderDataset"},"widgets_values":["/container/transfer/datasets/dog_and_cat/dog_and_cat/train/"]},{"id":56,"type":"MarkdownNote","pos":[3464.170166015625,1519.1202392578125],"size":[285.6708679199219,95.45172882080078],"flags":{},"order":15,"mode":2,"inputs":[],"outputs":[],"properties":{},"widgets_values":["# Save model\nUpdate the model path if needed.\nThis is saved under custom_nodes/data-analysis/models"],"color":"#432","bgcolor":"#653"}],"links":[[68,11,0,50,2,"PTOPTIMIZER"],[70,19,0,51,0,"STRING"],[71,50,0,19,0,"PTMODEL"],[72,50,1,57,0,"TENSOR"],[73,50,2,58,0,"TENSOR"],[74,57,0,59,0,"DATAFRAME"],[75,58,0,59,1,"DATAFRAME"],[77,64,0,65,1,"DATAFRAME"],[78,62,0,64,0,"TENSOR"],[79,59,0,65,0,"DATAFRAME"],[80,65,0,66,0,"DATAFRAME"],[81,65,0,61,0,"DATAFRAME"],[82,61,0,60,0,"DATAFRAME"],[84,60,0,68,0,"IMAGE"],[86,69,0,72,0,"PTVDATASET"],[87,72,0,50,1,"PTDATALOADER"],[88,71,0,50,0,"PTMODEL"],[89,71,0,11,0,"PTMODEL"],[91,74,0,75,0,"PTVDATASET"],[93,75,0,50,3,"PTDATALOADER"],[94,78,0,69,0,"PTVTRANSFORM"],[95,78,0,74,0,"PTVTRANSFORM"]],"groups":[{"id":1,"title":"Training metrics","bounding":[3447.499755859375,2242.806884765625,2284.0107421875,790.1612548828125],"color":"#8A8","font_size":24,"flags":{}}],"config":{},"extra":{"ds":{"scale":0.3736662716119983,"offset":[-975.7677617863449,-688.047741909445]},"node_versions":{"data-analysis":"db92cda5f5a1e611e15ca6f2531ff8585a5607a5","comfy-core":"v0.3.10-40-gd0f3752"},"ue_links":[]},"version":0.4} \ No newline at end of file diff --git a/modules/pytorch_wrapper/ptv_transforms_resize.py b/modules/pytorch_wrapper/ptv_transforms_resize.py new file mode 100644 index 0000000..39b5eb7 --- /dev/null +++ b/modules/pytorch_wrapper/ptv_transforms_resize.py @@ -0,0 +1,48 @@ +from typing import Any, Dict +import torchvision.transforms as transforms + + +class PtvTransformsResize: + """ + Resizes and transforms elements of dataset to PyTorch tensors. + + category: PyTorch wrapper - Training + """ + + @classmethod + def INPUT_TYPES(cls) -> Dict[str, Any]: + """ + Defines the required input types. + + Returns: + Dict[str, Any]: An empty dictionary. + """ + return { + "required": { + "height": ("INT", {"default": 256, "min": 1, "max": 32768}), + "width": ("INT", {"default": 256, "min": 1, "max": 32768}) + }, + } + + RETURN_TYPES: tuple = ("PTVTRANSFORM",) + FUNCTION: str = "f" + CATEGORY: str = "Training" + + def f(self, height: int, width: int) -> tuple: + """ + Resizes and transforms elements of dataset to PyTorch tensors. + + Args: + height (int): Target height of the input image after transform + width (int): Target width of the inpt image after transform + + Returns: + tuple: A tuple containing the transform callable. + """ + c = transforms.Compose([ + transforms.Resize((height, width)), + transforms.ToTensor() + ]) + return(c,) + + diff --git a/pyproject.toml b/pyproject.toml index d73233a..8152dc7 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.7" +version = "2.0.8" 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 index bfb5e7a..cf6419e 100644 --- a/tests/pytorch_wrapper/pt_evaluate_classification_conv_model_test2.py +++ b/tests/pytorch_wrapper/pt_evaluate_classification_conv_model_test2.py @@ -9,17 +9,19 @@ 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.ptv_transforms_resize import PtvTransformsResize 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): + + @unittest.skipIf(os.getenv("RUN_SKIPPED_TESTS") != "1", "Skipping unless explicitly enabled") def setUp(self): """Set up test instance.""" # Transform - compose_node = PtvTransformsToTensor() - transform = compose_node.f()[0] + compose_node = PtvTransformsResize() + transform = compose_node.f(256, 256)[0] # Dataset self.node = PtvImageFolderDataset() @@ -38,10 +40,10 @@ class TestPtEvaluate(unittest.TestCase): # Model self.model_node = PtnConvModel() self.model = self.model_node.f( - input_dim="(3, 512, 512)", + input_dim="(3, 256, 256)", penultimate_dim=0, output_dim=2, - channel_list="[32,64,128,256,512]", # To 256, 128, 64, 32, 16 + channel_list="[32,64,128,256,512]", # To 128, 64, 32, 16, 8 kernel_size_list="[3,3,3,3,3]", padding_list="[1,1,1,1,1]", downsample_list="[True,True,True,True,True]" @@ -49,12 +51,13 @@ class TestPtEvaluate(unittest.TestCase): # Load weights self.load_model = PtLoadModel() - epochs = 7 + epochs = 20 self.loaded_model = self.load_model.f(self.model,f"dog_cat_{epochs}_epochs_conv.pt") # Instantiates evaluation node self.evaluate = PtEvaluateClassificationModel() + @unittest.skipIf(os.getenv("RUN_SKIPPED_TESTS") != "1", "Skipping unless explicitly enabled") def test_evaluation(self): """ Tests evaluation. diff --git a/tests/pytorch_wrapper/pt_train_classification_conv_model_test2.py b/tests/pytorch_wrapper/pt_train_classification_conv_model_test2.py index b6fc7d0..1c38ad6 100644 --- a/tests/pytorch_wrapper/pt_train_classification_conv_model_test2.py +++ b/tests/pytorch_wrapper/pt_train_classification_conv_model_test2.py @@ -11,23 +11,34 @@ 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.ptv_transforms_resize import PtvTransformsResize 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): + + @unittest.skipIf(os.getenv("RUN_SKIPPED_TESTS") != "1", "Skipping unless explicitly enabled") def setUp(self): """Set up test instance.""" - compose_node = PtvTransformsToTensor() - transform = compose_node.f()[0] + compose_node = PtvTransformsResize() + transform = compose_node.f(256, 256)[0] self.node = PtvImageFolderDataset() - dataset = self.node.f("dog_and_cat/train", + train_dataset = self.node.f("dog_and_cat/train", transform=transform)[0] - data_loader_node = PtDataLoader() - self.train_loader = data_loader_node.f( - dataset=dataset, + train_data_loader_node = PtDataLoader() + val_dataset = self.node.f("dog_and_cat/val", + transform=transform)[0] + val_data_loader_node = PtDataLoader() + self.train_loader = train_data_loader_node.f( + dataset=train_dataset, + batch_size=32, + shuffle=True, + parameters='{"num_workers":1}' + )[0] + self.val_loader = val_data_loader_node.f( + dataset=val_dataset, batch_size=32, shuffle=True, parameters='{"num_workers":1}' @@ -35,10 +46,10 @@ class TestPtTrainConvModel2(unittest.TestCase): self.model_node = PtnConvModel() self.model = self.model_node.f( - input_dim="(3, 512, 512)", + input_dim="(3, 256, 256)", penultimate_dim=0, output_dim=2, - channel_list="[32,64,128,256,512]", # To 256, 128, 64, 32, 16 + channel_list="[32,64,128,256,512]", # To 128, 64, 32, 16, 8 kernel_size_list="[3,3,3,3,3]", padding_list="[1,1,1,1,1]", downsample_list="[True,True,True,True,True]" @@ -48,11 +59,13 @@ class TestPtTrainConvModel2(unittest.TestCase): self.trainer = PtTrainClassificationModel() self.save_model = PtSaveModel() + @unittest.skipIf(os.getenv("RUN_SKIPPED_TESTS") != "1", "Skipping unless explicitly enabled") def test_1(self): - epochs = 7 + epochs = 20 trained_model = self.trainer.f(self.model, self.train_loader, self.optimizer, epochs, # epochs - use_gpu=True)[0] + use_gpu=True, + val_loader=self.val_loader)[0] self.save_model.f(trained_model, f"dog_cat_{epochs}_epochs_conv.pt") diff --git a/tests/pytorch_wrapper/ptv_transforms_resize_test.py b/tests/pytorch_wrapper/ptv_transforms_resize_test.py new file mode 100644 index 0000000..ca6d386 --- /dev/null +++ b/tests/pytorch_wrapper/ptv_transforms_resize_test.py @@ -0,0 +1,43 @@ +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_resize import PtvTransformsResize + + +class TestPtvTransformsResize(unittest.TestCase): + def setUp(self): + """Set up test data.""" + self.compose_node = PtvTransformsResize() + self.height = 128 + self.width = 128 + + def test_with_transform(self): + """Tests the transformation output.""" + transform = self.compose_node.f(self.height, self.width)[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, self.height, self.width), "Output shape should match target dimensions.") + + def test_default_parameters(self): + """Tests default transform parameters.""" + transform = self.compose_node.f(256, 256)[0] + self.assertTrue(callable(transform), "Default transform should be callable.") + + +if __name__ == "__main__": + unittest.main()