nodes/HowToSD-ComfyUI-Data-Analysis
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Added image resizing related nodes and a single image prediction node

Browse Source
This commit is contained in:
HowToSD
2025-03-08 03:38:52 -08:00
parent a2fc03f41f
commit 5fee9ffb60
32 changed files with 1007 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files
__init__.py
+21
View File
@@ -189,6 +189,7 @@ from .modules.pytorch_wrapper.pt_bmm import PtBmm
from .modules.pytorch_wrapper.pt_bool_create import PtBoolCreate
from .modules.pytorch_wrapper.pt_cos import PtCos
from .modules.pytorch_wrapper.pt_cosh import PtCosh
from .modules.pytorch_wrapper.pt_crop import PtCrop
from .modules.pytorch_wrapper.pt_data_loader import PtDataLoader
from .modules.pytorch_wrapper.pt_div import PtDiv
from .modules.pytorch_wrapper.pt_einsum import PtEinsum
@@ -199,6 +200,7 @@ from .modules.pytorch_wrapper.pt_flatten import PtFlatten
from .modules.pytorch_wrapper.pt_float_create import PtFloatCreate
from .modules.pytorch_wrapper.pt_floor_divide import PtFloorDiv
from .modules.pytorch_wrapper.pt_from_image import PtFromImage
from .modules.pytorch_wrapper.pt_from_image_transpose import PtFromImageTranspose
from .modules.pytorch_wrapper.pt_from_latent import PtFromLatent
from .modules.pytorch_wrapper.pt_from_numpy import PtFromNumpy
from .modules.pytorch_wrapper.pt_full import PtFull
@@ -207,6 +209,8 @@ from .modules.pytorch_wrapper.pt_ge import PtGe
from .modules.pytorch_wrapper.pt_gt import PtGt
from .modules.pytorch_wrapper.pt_index_select import PtIndexSelect
from .modules.pytorch_wrapper.pt_int_create import PtIntCreate
from .modules.pytorch_wrapper.pt_interpolate_by_scale_factor import PtInterpolateByScaleFactor
from .modules.pytorch_wrapper.pt_interpolate_to_size import PtInterpolateToSize
from .modules.pytorch_wrapper.pt_le import PtLe
from .modules.pytorch_wrapper.pt_linspace import PtLinspace
from .modules.pytorch_wrapper.pt_load_model import PtLoadModel
@@ -227,8 +231,10 @@ from .modules.pytorch_wrapper.pt_mul import PtMul
from .modules.pytorch_wrapper.pt_ne import PtNe
from .modules.pytorch_wrapper.pt_neg import PtNeg
from .modules.pytorch_wrapper.pt_ones import PtOnes
from .modules.pytorch_wrapper.pt_pad import PtPad
from .modules.pytorch_wrapper.pt_permute import PtPermute
from .modules.pytorch_wrapper.pt_pow import PtPow
from .modules.pytorch_wrapper.pt_predict_classification_model import PtPredictClassificationModel
from .modules.pytorch_wrapper.pt_prod import PtProd
from .modules.pytorch_wrapper.pt_rand import PtRand
from .modules.pytorch_wrapper.pt_rand_int import PtRandInt
@@ -253,6 +259,7 @@ from .modules.pytorch_wrapper.pt_sum import PtSum
from .modules.pytorch_wrapper.pt_tan import PtTan
from .modules.pytorch_wrapper.pt_tanh import PtTanh
from .modules.pytorch_wrapper.pt_to_image import PtToImage
from .modules.pytorch_wrapper.pt_to_image_transpose import PtToImageTranspose
from .modules.pytorch_wrapper.pt_to_latent import PtToLatent
from .modules.pytorch_wrapper.pt_to_numpy import PtToNumpy
from .modules.pytorch_wrapper.pt_to_rgb_tensors import PtToRgbTensors
@@ -464,6 +471,7 @@ NODE_CLASS_MAPPINGS: Dict[str, Type[T]] = {
"PtBoolCreate": PtBoolCreate,
"PtCos": PtCos,
"PtCosh": PtCosh,
"PtCrop": PtCrop,
"PtDataLoader": PtDataLoader,
"PtDiv": PtDiv,
"PtEinsum": PtEinsum,
@@ -474,6 +482,7 @@ NODE_CLASS_MAPPINGS: Dict[str, Type[T]] = {
"PtFloatCreate": PtFloatCreate,
"PtFloorDiv": PtFloorDiv,
"PtFromImage": PtFromImage,
"PtFromImageTranspose": PtFromImageTranspose,
"PtFromLatent": PtFromLatent,
"PtFromNumpy": PtFromNumpy,
"PtFull": PtFull,
@@ -482,6 +491,8 @@ NODE_CLASS_MAPPINGS: Dict[str, Type[T]] = {
"PtGt": PtGt,
"PtIndexSelect": PtIndexSelect,
"PtIntCreate": PtIntCreate,
"PtInterpolateByScaleFactor": PtInterpolateByScaleFactor,
"PtInterpolateToSize": PtInterpolateToSize,
"PtLe": PtLe,
"PtLinspace": PtLinspace,
"PtLoadModel": PtLoadModel,
@@ -502,8 +513,10 @@ NODE_CLASS_MAPPINGS: Dict[str, Type[T]] = {
"PtNe": PtNe,
"PtNeg": PtNeg,
"PtOnes": PtOnes,
"PtPad": PtPad,
"PtPermute": PtPermute,
"PtPow": PtPow,
"PtPredictClassificationModel": PtPredictClassificationModel,
"PtProd": PtProd,
"PtRand": PtRand,
"PtRandInt": PtRandInt,
@@ -528,6 +541,7 @@ NODE_CLASS_MAPPINGS: Dict[str, Type[T]] = {
"PtTan": PtTan,
"PtTanh": PtTanh,
"PtToImage": PtToImage,
"PtToImageTranspose": PtToImageTranspose,
"PtToLatent": PtToLatent,
"PtToNumpy": PtToNumpy,
"PtToRgbTensors": PtToRgbTensors,
@@ -750,6 +764,7 @@ NODE_DISPLAY_NAME_MAPPINGS: Dict[str, str] = {
"PtBoolCreate": "Pt Bool Create",
"PtCos": "Pt Cos",
"PtCosh": "Pt Cosh",
"PtCrop": "Pt Crop",
"PtDataLoader": "Pt Data Loader",
"PtDiv": "Pt Div",
"PtEinsum": "Pt Einsum",
@@ -760,6 +775,7 @@ NODE_DISPLAY_NAME_MAPPINGS: Dict[str, str] = {
"PtFloatCreate": "Pt Float Create",
"PtFloorDiv": "Pt Floor Div",
"PtFromImage": "Pt From Image",
"PtFromImageTranspose": "Pt From Image Transpose",
"PtFromLatent": "Pt From Latent",
"PtFromNumpy": "Pt From Numpy",
"PtFull": "Pt Full",
@@ -768,6 +784,8 @@ NODE_DISPLAY_NAME_MAPPINGS: Dict[str, str] = {
"PtGt": "Pt Gt",
"PtIndexSelect": "Pt Index Select",
"PtIntCreate": "Pt Int Create",
"PtInterpolateByScaleFactor": "Pt Interpolate By Scale Factor",
"PtInterpolateToSize": "Pt Interpolate To Size",
"PtLe": "Pt Le",
"PtLinspace": "Pt Linspace",
"PtLoadModel": "Pt Load Model",
@@ -788,8 +806,10 @@ NODE_DISPLAY_NAME_MAPPINGS: Dict[str, str] = {
"PtNe": "Pt Ne",
"PtNeg": "Pt Neg",
"PtOnes": "Pt Ones",
"PtPad": "Pt Pad",
"PtPermute": "Pt Permute",
"PtPow": "Pt Pow",
"PtPredictClassificationModel": "Pt Predict Classification Model",
"PtProd": "Pt Prod",
"PtRand": "Pt Rand",
"PtRandInt": "Pt Rand Int",
@@ -814,6 +834,7 @@ NODE_DISPLAY_NAME_MAPPINGS: Dict[str, str] = {
"PtTan": "Pt Tan",
"PtTanh": "Pt Tanh",
"PtToImage": "Pt To Image",
"PtToImageTranspose": "Pt To Image Transpose",
"PtToLatent": "Pt To Latent",
"PtToNumpy": "Pt To Numpy",
"PtToRgbTensors": "Pt To Rgb Tensors",
docs/reference/node_reference.md
+13 -3
View File
@@ -206,6 +206,17 @@
| [Pt Le](pt_le.md) | Tests whether elements in the first PyTorch tensor are less than or equal to the corresponding elements in the second tensor. |
| [Pt Lt](pt_lt.md) | Tests whether elements in the first PyTorch tensor are less than the corresponding elements in the second tensor. |
| [Pt Ne](pt_ne.md) | Tests whether two PyTorch tensors are not equal element-wise. |
## PyTorch wrapper - Image processing
| Node | Description |
| --- | --- |
| [Pt Crop](pt_crop.md) | Crops a PyTorch tensor to the specified size. The input tensor must have a shape of (c, h, w) or (b, c, h, w). |
| [Pt From Image](pt_from_image.md) | Casts an Image tensor as a PyTorch tensor. |
| [Pt From Image Transpose](pt_from_image_transpose.md) | Casts an image tensor to a PyTorch tensor and transposes it from (H, W, C) to (C, H, W). For rank-4 inputs, the batch axis remains unchanged. |
| [Pt Interpolate By Scale Factor](pt_interpolate_by_scale_factor.md) | Resizes a PyTorch tensor using interpolation by scale factor. The input tensor must have a shape of (c, h, w) or (b, c, h, w). |
| [Pt Interpolate To Size](pt_interpolate_to_size.md) | Resizes a PyTorch tensor using interpolation. The input tensor must have a shape of (c, h, w) or (b, c, h, w). |
| [Pt Pad](pt_pad.md) | Pads a PyTorch tensor to the specified size. Padded area will be black. The input tensor must have a shape of (c, h, w) or (b, c, h, w). |
| [Pt To Image](pt_to_image.md) | Casts a PyTorch tensor as an Image tensor. |
| [Pt To Image Transpose](pt_to_image_transpose.md) | Casts a PyTorch tensor as an Image tensor and transposes it from (C, H, W) to (H, W, C). For rank-4 inputs, the batch axis remains unchanged. |
## PyTorch wrapper - Indexing and Slicing Operations
| Node | Description |
| --- | --- |
@@ -271,7 +282,6 @@
| [Pt Arange](pt_arange.md) | Creates a PyTorch tensor using `torch.arange` with the specified start, end, and step values. |
| [Pt Bool Create](pt_bool_create.md) | Creates a PyTorch tensor of dtype bool from True or False values entered as a list in the text field. |
| [Pt Float Create](pt_float_create.md) | Creates a PyTorch tensor with 32-bit floating point precision |
| [Pt From Image](pt_from_image.md) | Casts an Image tensor as a PyTorch tensor. |
| [Pt From Latent](pt_from_latent.md) | Casts a latent tensor as a PyTorch tensor. |
| [Pt From Numpy](pt_from_numpy.md) | Converts a NumPy ndarray to a PyTorch tensor while preserving its data type. |
| [Pt Full](pt_full.md) | Creates a PyTorch tensor filled with a specified value using the size entered in the text field. |
@@ -285,20 +295,20 @@
## PyTorch wrapper - Tensor data conversion
| Node | Description |
| --- | --- |
| [Pt To Image](pt_to_image.md) | Casts a PyTorch tensor as an Image tensor. |
| [Pt To Latent](pt_to_latent.md) | Casts a PyTorch tensor as a latent tensor. |
| [Pt To Numpy](pt_to_numpy.md) | Converts PyTorch tensor to NumPy ndarray. |
| [Pt To Rgb Tensors](pt_to_rgb_tensors.md) | Splits a PyTorch tensor into R, G, and B tensors. |
| [Pto Adam](pto_adam.md) | Converts PyTorch tensor to NumPy ndarray. |
## PyTorch wrapper - Training
| Node | Description |
| --- | --- |
| [Pt Data Loader](pt_data_loader.md) | Loads data from a dataset node and creates a PyTorch DataLoader. |
| [Pt Evaluate Classification Model](pt_evaluate_classification_model.md) | Performs inference on test data and computes evaluation metrics. |
| [Pt Load Model](pt_load_model.md) | A wrapper class for saving a PyTorch model. |
| [Pt Predict Classification Model](pt_predict_classification_model.md) | Performs inference on input data. |
| [Pt Train Classification Model](pt_train_classification_model.md) | Trains a classification model using a given dataset, optimizer, and number of epochs. |
| [Ptn Conv Model](ptn_conv_model.md) | A convolutional model consisting of multiple convolutional layers. |
| [Ptn Linear Model](ptn_linear_model.md) | A linear model consisting of dense layers. |
| [Pto Adam](pto_adam.md) | Instantiates the Adam optimizer. |
| [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. |
docs/reference/pt_crop.md
+23
View File
@@ -0,0 +1,23 @@
# Pt Crop
Crops a PyTorch tensor to the specified size. The input tensor must have a shape of (c, h, w) or (b, c, h, w).
In ComfyUI Analysis, the tensor data type is "TENSOR," while in ComfyUI, the type "IMAGE" is used. To pass image data from ComfyUI to a ComfyUI Analysis node (e.g., this interpolation node), first use the "Pt From Image" node to convert it to a tensor, then transpose the axes from (b, h, w, c) to (b, c, h, w) before passing it to this node.
This transposition can be performed using the "Pt Permute" node with (0, 2, 3, 1). To convert the output back to an image, apply "Pt Permute" with (0, 3, 1, 2) and then use the "Pt To Image" node.
## Input
| Name | Data type |
|---|---|
| tens | Tensor |
| height | Int |
| width | Int |
## Output
| Data type |
|---|
| Tensor |
<HR>
Category: PyTorch wrapper - Image processing
ComfyUI Data Analysis Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved.
docs/reference/pt_from_image.md
+1 -1
View File
@@ -12,6 +12,6 @@ Casts an Image tensor as a PyTorch tensor.
| Tensor |
<HR>
Category: PyTorch wrapper - Tensor creation
Category: PyTorch wrapper - Image processing
ComfyUI Data Analysis Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved.
docs/reference/pt_from_image_transpose.md
+17
View File
@@ -0,0 +1,17 @@
# Pt From Image Transpose
Casts an image tensor to a PyTorch tensor and transposes it from (H, W, C) to (C, H, W). For rank-4 inputs, the batch axis remains unchanged.
## Input
| Name | Data type |
|---|---|
| image | Image |
## Output
| Data type |
|---|
| Tensor |
<HR>
Category: PyTorch wrapper - Image processing
ComfyUI Data Analysis Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved.
docs/reference/pt_interpolate_by_scale_factor.md
+23
View File
@@ -0,0 +1,23 @@
# Pt Interpolate By Scale Factor
Resizes a PyTorch tensor using interpolation by scale factor. The input tensor must have a shape of (c, h, w) or (b, c, h, w).
In ComfyUI Analysis, the tensor data type is "TENSOR," while in ComfyUI, the type "IMAGE" is used. To pass image data from ComfyUI to a ComfyUI Analysis node (e.g., this interpolation node), first use the "Pt From Image" node to convert it to a tensor, then transpose the axes from (b, h, w, c) to (b, c, h, w) before passing it to this node.
This transposition can be performed using the "Pt Permute" node with (0, 2, 3, 1). To convert the output back to an image, apply "Pt Permute" with (0, 3, 1, 2) and then use the "Pt To Image" node.
## Input
| Name | Data type |
|---|---|
| tens | Tensor |
| scale_factor | Float |
| mode | |
## Output
| Data type |
|---|
| Tensor |
<HR>
Category: PyTorch wrapper - Image processing
ComfyUI Data Analysis Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved.
docs/reference/pt_interpolate_to_size.md
+24
View File
@@ -0,0 +1,24 @@
# Pt Interpolate To Size
Resizes a PyTorch tensor using interpolation. The input tensor must have a shape of (c, h, w) or (b, c, h, w).
In ComfyUI Analysis, the tensor data type is "TENSOR," while in ComfyUI, the type "IMAGE" is used. To pass image data from ComfyUI to a ComfyUI Analysis node (e.g., this interpolation node), first use the "Pt From Image" node to convert it to a tensor, then transpose the axes from (b, h, w, c) to (b, c, h, w) before passing it to this node.
This transposition can be performed using the "Pt Permute" node with (0, 2, 3, 1). To convert the output back to an image, apply "Pt Permute" with (0, 3, 1, 2) and then use the "Pt To Image" node.
## Input
| Name | Data type |
|---|---|
| tens | Tensor |
| height | Int |
| width | Int |
| mode | |
## Output
| Data type |
|---|
| Tensor |
<HR>
Category: PyTorch wrapper - Image processing
ComfyUI Data Analysis Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved.
docs/reference/pt_pad.md
+23
View File
@@ -0,0 +1,23 @@
# Pt Pad
Pads a PyTorch tensor to the specified size. Padded area will be black. The input tensor must have a shape of (c, h, w) or (b, c, h, w).
In ComfyUI Analysis, the tensor data type is "TENSOR," while in ComfyUI, the type "IMAGE" is used. To pass image data from ComfyUI to a ComfyUI Analysis node (e.g., this interpolation node), first use the "Pt From Image" node to convert it to a tensor, then transpose the axes from (b, h, w, c) to (b, c, h, w) before passing it to this node.
This transposition can be performed using the "Pt Permute" node with (0, 2, 3, 1). To convert the output back to an image, apply "Pt Permute" with (0, 3, 1, 2) and then use the "Pt To Image" node.
## Input
| Name | Data type |
|---|---|
| tens | Tensor |
| height | Int |
| width | Int |
## Output
| Data type |
|---|
| Tensor |
<HR>
Category: PyTorch wrapper - Image processing
ComfyUI Data Analysis Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved.
docs/reference/pt_predict_classification_model.md
+22
View File
@@ -0,0 +1,22 @@
# Pt Predict Classification Model
Performs inference on input data.
## Input
| Name | Data type |
|---|---|
| model | Ptmodel |
| inputs | Tensor |
| class_id_to_name_map | Pydict |
| use_gpu | Boolean |
## Output
| Data type |
|---|
| String |
| Int |
| Float |
<HR>
Category: PyTorch wrapper - Training
ComfyUI Data Analysis Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved.
docs/reference/pt_to_image.md
+1 -1
View File
@@ -12,6 +12,6 @@ Casts a PyTorch tensor as an Image tensor.
| Image |
<HR>
Category: PyTorch wrapper - Tensor data conversion
Category: PyTorch wrapper - Image processing
ComfyUI Data Analysis Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved.
docs/reference/pt_to_image_transpose.md
+17
View File
@@ -0,0 +1,17 @@
# Pt To Image Transpose
Casts a PyTorch tensor as an Image tensor and transposes it from (C, H, W) to (H, W, C). For rank-4 inputs, the batch axis remains unchanged.
## Input
| Name | Data type |
|---|---|
| tens | Tensor |
## Output
| Data type |
|---|
| Image |
<HR>
Category: PyTorch wrapper - Image processing
ComfyUI Data Analysis Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved.
docs/reference/pto_adam.md
+2 -2
View File
@@ -1,5 +1,5 @@
# Pto Adam
Converts PyTorch tensor to NumPy ndarray.
Instantiates the Adam optimizer.
## Input
| Name | Data type |
@@ -15,6 +15,6 @@ Converts PyTorch tensor to NumPy ndarray.
| Ptoptimizer |
<HR>
Category: PyTorch wrapper - Tensor data conversion
Category: PyTorch wrapper - Training
ComfyUI Data Analysis Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved.
examples/workflows/dog_cat_classification_prediction.json
+1
View File
@@ -0,0 +1 @@
{"last_node_id":74,"last_link_id":97,"nodes":[{"id":73,"type":"PtFromImageTranspose","pos":[1488.2567138671875,1628.712646484375],"size":[289.79998779296875,26],"flags":{},"order":4,"mode":0,"inputs":[{"name":"image","type":"IMAGE","link":95}],"outputs":[{"name":"TENSOR","type":"TENSOR","links":[96],"slot_index":0}],"properties":{"Node name for S&R":"PtFromImageTranspose"}},{"id":74,"type":"PtInterpolateToSize","pos":[1801.049560546875,1628.712646484375],"size":[315,106],"flags":{},"order":5,"mode":0,"inputs":[{"name":"tens","type":"TENSOR","link":96}],"outputs":[{"name":"TENSOR","type":"TENSOR","links":[97],"slot_index":0}],"properties":{"Node name for S&R":"PtInterpolateToSize"},"widgets_values":[256,256,"nearest-exact"]},{"id":69,"type":"PtSqueeze","pos":[2206.266357421875,1639.8197021484375],"size":[210,60.5162467956543],"flags":{},"order":6,"mode":0,"inputs":[{"name":"tens","type":"TENSOR","link":97}],"outputs":[{"name":"TENSOR","type":"TENSOR","links":[88],"slot_index":0}],"properties":{"Node name for S&R":"PtSqueeze"},"widgets_values":[0]},{"id":55,"type":"PtPredictClassificationModel","pos":[2482.278564453125,1626.3341064453125],"size":[405.5999755859375,98],"flags":{},"order":7,"mode":0,"inputs":[{"name":"model","type":"PTMODEL","link":70},{"name":"inputs","type":"TENSOR","link":88},{"name":"class_id_to_name_map","type":"PYDICT","link":69}],"outputs":[{"name":"class_name","type":"STRING","links":[74],"slot_index":0},{"name":"class_id","type":"INT","links":null},{"name":"probability","type":"FLOAT","links":[93],"slot_index":2}],"properties":{"Node name for S&R":"PtPredictClassificationModel"},"widgets_values":[true]},{"id":56,"type":"PyStringToDict","pos":[2206.266357421875,1774.539306640625],"size":[210,81.30076599121094],"flags":{},"order":0,"mode":0,"inputs":[],"outputs":[{"name":"PYDICT","type":"PYDICT","links":[69],"slot_index":0}],"properties":{"Node name for S&R":"PyStringToDict"},"widgets_values":["{0:\"cat\", 1:\"dog\"}"]},{"id":60,"type":"PandasShowText","pos":[2970.7451171875,1629.65576171875],"size":[221.17684936523438,76],"flags":{},"order":8,"mode":0,"inputs":[{"name":"text","type":"STRING","link":74,"widget":{"name":"text"}}],"outputs":[{"name":"STRING","type":"STRING","links":null,"shape":6}],"properties":{"Node name for S&R":"PandasShowText"},"widgets_values":["","cat"]},{"id":72,"type":"CDAShowFloat","pos":[2976.763916015625,1770.6387939453125],"size":[210,76],"flags":{},"order":9,"mode":0,"inputs":[{"name":"float_scalar","type":"FLOAT","link":93,"widget":{"name":"float_scalar"}}],"outputs":[{"name":"STRING","type":"STRING","links":null,"shape":6}],"properties":{"Node name for S&R":"CDAShowFloat"},"widgets_values":[0,"0.9118371605873108"]},{"id":23,"type":"PtLoadModel","pos":[2206.266357421875,1482.663330078125],"size":[210,75.45188903808594],"flags":{},"order":3,"mode":0,"inputs":[{"name":"model","type":"PTMODEL","link":62}],"outputs":[{"name":"model","type":"PTMODEL","links":[70],"slot_index":0}],"properties":{"Node name for S&R":"PtLoadModel"},"widgets_values":["dog_cat.pt"]},{"id":57,"type":"LoadImage","pos":[1135.1787109375,1628.712646484375],"size":[315,314],"flags":{},"order":2,"mode":0,"inputs":[],"outputs":[{"name":"IMAGE","type":"IMAGE","links":[95],"slot_index":0},{"name":"MASK","type":"MASK","links":null}],"properties":{"Node name for S&R":"LoadImage"},"widgets_values":["cat.jpg","image"]},{"id":48,"type":"PtnConvModel","pos":[1801.049560546875,1306.3406982421875],"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]"]}],"links":[[62,48,0,23,0,"PTMODEL"],[69,56,0,55,2,"PYDICT"],[70,23,0,55,0,"PTMODEL"],[74,55,0,60,0,"STRING"],[88,69,0,55,1,"TENSOR"],[93,55,2,72,0,"FLOAT"],[95,57,0,73,0,"IMAGE"],[96,73,0,74,0,"TENSOR"],[97,74,0,69,0,"TENSOR"]],"groups":[],"config":{},"extra":{"ds":{"scale":0.6339486400897979,"offset":[-879.509850098028,-1085.7559202042573]},"node_versions":{"data-analysis":"fd6fa94b931cd094ce5766852f5914d7291cb464","comfy-core":"v0.3.10-40-gd0f3752"},"ue_links":[]},"version":0.4}
modules/pytorch_wrapper/pt_crop.py
+65
View File
@@ -0,0 +1,65 @@
from typing import Any, Dict
import torch
class PtCrop:
"""
Crops a PyTorch tensor to the specified size. The input tensor must have a shape of (c, h, w) or (b, c, h, w).
In ComfyUI Analysis, the tensor data type is "TENSOR," while in ComfyUI, the type "IMAGE" is used. To pass image data from ComfyUI to a ComfyUI Analysis node (e.g., this interpolation node), first use the "Pt From Image" node to convert it to a tensor, then transpose the axes from (b, h, w, c) to (b, c, h, w) before passing it to this node.
This transposition can be performed using the "Pt Permute" node with (0, 2, 3, 1). To convert the output back to an image, apply "Pt Permute" with (0, 3, 1, 2) and then use the "Pt To Image" node.
category: PyTorch wrapper - Image processing
"""
@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": {
"tens": ("TENSOR", {}),
"height": ("INT", {"default":0, "min":1, "max": 16384}),
"width": ("INT", {"default":0, "min":1, "max": 16384})
}
}
RETURN_TYPES: tuple = ("TENSOR",)
FUNCTION: str = "f"
CATEGORY: str = "Data Analysis"
def f(self, tens: torch.Tensor, height: int, width: int) -> tuple:
"""
Crops a PyTorch tensor.
Args:
tens (torch.Tensor): A PyTorch tensor with the shape of (c, h, w) or (b, c, h, w).
height (int): New image height.
width (int): New image width.
Returns:
tuple: A tuple containing a transformed tensor`.
"""
dim = tens.dim()
if dim == 3: # c, h, w
image_height = tens.size(1)
image_width = tens.size(2)
elif dim == 4: # b, c, h, w
image_height = tens.size(2)
image_width = tens.size(3)
else:
raise ValueError("Only rank 3 or 4 tensors are supported for crop.")
if height > image_height or width > image_width:
raise ValueError("Specified crop size exceeds tensor dimensions.")
# Compute starting position for cropping
top = max(int((height - image_height) / 2), 0)
left = max(int((width - image_width) / 2), 0)
tens2 = tens[..., top:top+height, left:left+width]
return (tens2,)
modules/pytorch_wrapper/pt_from_image.py
+1 -1
View File
@@ -6,7 +6,7 @@ class PtFromImage:
"""
Casts an Image tensor as a PyTorch tensor.
category: PyTorch wrapper - Tensor creation
category: PyTorch wrapper - Image processing
"""
@classmethod
modules/pytorch_wrapper/pt_from_image_transpose.py
+47
View File
@@ -0,0 +1,47 @@
from typing import Any, Dict
import torch
class PtFromImageTranspose:
"""
Casts an image tensor to a PyTorch tensor and transposes it from (H, W, C) to (C, H, W). For rank-4 inputs, the batch axis remains unchanged.
category: PyTorch wrapper - Image processing
"""
@classmethod
def INPUT_TYPES(cls) -> Dict[str, Any]:
"""
Defines the input types for the function.
Returns:
Dict[str, Any]: A dictionary specifying required input types.
"""
return {
"required": {
"image": ("IMAGE", {}),
}
}
RETURN_TYPES: tuple = ("TENSOR",)
FUNCTION: str = "f"
CATEGORY: str = "Data Analysis"
def f(self, image: torch.Tensor) -> tuple:
"""
Casts an Image tensor as a PyTorch tensor.
Args:
image (torch.Tensor): Image tensor
Returns:
tuple: A tuple containing the PyTorch tensor.
"""
if image.dim() == 3:
tens = image.permute(2, 0, 1)
elif image.dim() == 4:
tens = image.permute(0, 3, 1, 2)
else:
raise ValueError("Only rank 3 or 4 tensors are supported.")
return (tens,)
modules/pytorch_wrapper/pt_interpolate_by_scale_factor.py
+56
View File
@@ -0,0 +1,56 @@
from typing import Any, Dict
import torch
class PtInterpolateByScaleFactor:
"""
Resizes a PyTorch tensor using interpolation by scale factor. The input tensor must have a shape of (c, h, w) or (b, c, h, w).
In ComfyUI Analysis, the tensor data type is "TENSOR," while in ComfyUI, the type "IMAGE" is used. To pass image data from ComfyUI to a ComfyUI Analysis node (e.g., this interpolation node), first use the "Pt From Image" node to convert it to a tensor, then transpose the axes from (b, h, w, c) to (b, c, h, w) before passing it to this node.
This transposition can be performed using the "Pt Permute" node with (0, 2, 3, 1). To convert the output back to an image, apply "Pt Permute" with (0, 3, 1, 2) and then use the "Pt To Image" node.
category: PyTorch wrapper - Image processing
"""
@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": {
"tens": ("TENSOR", {}),
"scale_factor": ("FLOAT", {"default":2.0, "min":0.01, "max": 100.0, "step":0.01}),
"mode": (("nearest-exact", "bilinear", "bicubic", "nearest"),)
}
}
RETURN_TYPES: tuple = ("TENSOR",)
FUNCTION: str = "f"
CATEGORY: str = "Data Analysis"
def f(self, tens: torch.Tensor, scale_factor: float, mode: str) -> tuple:
"""
Resizes a PyTorch tensor using interpolation.
Args:
tens (torch.Tensor): A PyTorch tensor with the shape of (c, h, w) or (b, c, h, w).
scale_factor (float): Scale factor.
mode (string): Interpolation mode.
Returns:
tuple: A tuple containing a transformed tensor`.
"""
if tens.dim() == 3:
rank3 = True
tens_adjusted = torch.unsqueeze(tens, 0)
else:
rank3 = False
tens_adjusted = tens
tens2 = torch.nn.functional.interpolate(tens_adjusted, scale_factor=scale_factor, mode=mode)
if rank3:
tens2 = torch.squeeze(tens2, 0)
return (tens2,)
modules/pytorch_wrapper/pt_interpolate_to_size.py
+60
View File
@@ -0,0 +1,60 @@
from typing import Any, Dict
import torch
class PtInterpolateToSize:
"""
Resizes a PyTorch tensor using interpolation. The input tensor must have a shape of (c, h, w) or (b, c, h, w).
In ComfyUI Analysis, the tensor data type is "TENSOR," while in ComfyUI, the type "IMAGE" is used. To pass image data from ComfyUI to a ComfyUI Analysis node (e.g., this interpolation node), first use the "Pt From Image" node to convert it to a tensor, then transpose the axes from (b, h, w, c) to (b, c, h, w) before passing it to this node.
This transposition can be performed using the "Pt Permute" node with (0, 2, 3, 1). To convert the output back to an image, apply "Pt Permute" with (0, 3, 1, 2) and then use the "Pt To Image" node.
category: PyTorch wrapper - Image processing
"""
@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": {
"tens": ("TENSOR", {}),
"height": ("INT", {"default":0, "min":1, "max": 16384}),
"width": ("INT", {"default":0, "min":1, "max": 16384}),
"mode": (("nearest-exact", "bilinear", "bicubic", "nearest"),)
}
}
RETURN_TYPES: tuple = ("TENSOR",)
FUNCTION: str = "f"
CATEGORY: str = "Data Analysis"
def f(self, tens: torch.Tensor, height: int, width: int, mode: str) -> tuple:
"""
Resizes a PyTorch tensor using interpolation.
Args:
tens (torch.Tensor): A PyTorch tensor with the shape of (c, h, w) or (b, c, h, w).
height (int): New image height.
width (int): New image width.
mode (string): Interpolation mode.
Returns:
tuple: A tuple containing a transformed tensor`.
"""
if tens.dim() == 3:
rank3 = True
tens_adjusted = torch.unsqueeze(tens, 0)
else:
rank3 = False
tens_adjusted = tens
tens2 = torch.nn.functional.interpolate(tens_adjusted, size=(height, width), mode=mode)
if rank3:
tens2 = torch.squeeze(tens2, 0)
return (tens2,)
modules/pytorch_wrapper/pt_load_model.py
+1 -1
View File
@@ -36,6 +36,6 @@ class PtLoadModel:
def f(self, model:Any, model_path: str):
model_full_path = get_model_full_path(model_path)
sd = torch.load(model_full_path)
sd = torch.load(model_full_path, weights_only=True)
model.load_state_dict(sd)
return (model,)
modules/pytorch_wrapper/pt_pad.py
+72
View File
@@ -0,0 +1,72 @@
from typing import Any, Dict
import torch
class PtPad:
"""
Pads a PyTorch tensor to the specified size. Padded area will be black. The input tensor must have a shape of (c, h, w) or (b, c, h, w).
In ComfyUI Analysis, the tensor data type is "TENSOR," while in ComfyUI, the type "IMAGE" is used. To pass image data from ComfyUI to a ComfyUI Analysis node (e.g., this interpolation node), first use the "Pt From Image" node to convert it to a tensor, then transpose the axes from (b, h, w, c) to (b, c, h, w) before passing it to this node.
This transposition can be performed using the "Pt Permute" node with (0, 2, 3, 1). To convert the output back to an image, apply "Pt Permute" with (0, 3, 1, 2) and then use the "Pt To Image" node.
category: PyTorch wrapper - Image processing
"""
@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": {
"tens": ("TENSOR", {}),
"height": ("INT", {"default":0, "min":1, "max": 16384}),
"width": ("INT", {"default":0, "min":1, "max": 16384})
}
}
RETURN_TYPES: tuple = ("TENSOR",)
FUNCTION: str = "f"
CATEGORY: str = "Data Analysis"
def f(self, tens: torch.Tensor, height: int, width: int) -> tuple:
"""
Pads a PyTorch tensor.
Args:
tens (torch.Tensor): A PyTorch tensor with the shape of (c, h, w) or (b, c, h, w).
height (int): New image height.
width (int): New image width.
Returns:
tuple: A tuple containing a transformed tensor.
"""
dim = tens.dim()
if dim == 3: # c, h, w
image_height = tens.size(1)
image_width = tens.size(2)
elif dim == 4: # b, c, h, w
image_height = tens.size(2)
image_width = tens.size(3)
else:
raise ValueError("Only rank 3 or 4 tensors are supported for padding.")
if height < image_height or width < image_width:
raise ValueError("Specified dimensions must be greater than or equal to tensor dimensions.")
# Compute starting position for padding
top = max((height - image_height) // 2, 0)
left = max((width - image_width) // 2, 0)
if dim == 3:
tens2 = torch.zeros(tens.size(0), height, width,
dtype=tens.dtype, device=tens.device)
else:
tens2 = torch.zeros(tens.size(0), tens.size(1), height, width,
dtype=tens.dtype, device=tens.device)
tens2[..., top:top+image_height, left:left+image_width] = tens
return (tens2,)
modules/pytorch_wrapper/pt_predict_classification_model.py
+66
View File
@@ -0,0 +1,66 @@
from typing import Any, Dict, Tuple
import torch
import torch.nn as nn
import torch.nn.functional as F
from sklearn.metrics import accuracy_score, precision_score, recall_score, f1_score
class PtPredictClassificationModel:
"""
Performs inference on input data.
category: PyTorch wrapper - Training
"""
@classmethod
def INPUT_TYPES(cls) -> Dict[str, Any]:
"""
Defines the input types for the function.
Returns:
Dict[str, Any]: A dictionary specifying required input types.
"""
return {
"required": {
"model": ("PTMODEL", {}),
"inputs": ("TENSOR", {}),
"class_id_to_name_map": ("PYDICT", {}),
"use_gpu": ("BOOLEAN", {"default":False})
}
}
RETURN_NAMES: Tuple[str, ...] = ("class_name", "class_id", "probability")
RETURN_TYPES: Tuple[str, ...] = ("STRING", "INT", "FLOAT")
FUNCTION: str = "f"
CATEGORY: str = "Training"
def f(self, model: nn.Module,
inputs: torch.Tensor,
class_id_to_name_map: Dict,
use_gpu: bool) -> Tuple[str, int, float]:
"""
Performs inference on test data and computes evaluation metrics.
Args:
model (torch.nn.Module): The PyTorch model to evaluate.
inputs (torch.Tensor): Input image to predict the class for.
class_id_to_name_map (Dict): A Python dictionary contains the mapping of
class ID to class name.
use_gpu (bool): Whether to use GPU for inference.
Returns:
Tuple[str, int, float]: A tuple containing "class_name", "class_id", "probability".
"""
with torch.inference_mode():
if use_gpu:
model.to("cuda")
model.train(False) # TODO: Change to model.*e*v*a*l() once Comfy's security checker is fixed.
x = torch.unsqueeze(inputs, 0)
if use_gpu:
x = x.to("cuda")
outputs = model(x)
class_id = outputs.argmax(dim=1)[0].detach().cpu().item()
prob = F.softmax(outputs, dim=1)[0][class_id].detach().cpu().item()
class_name = class_id_to_name_map[class_id]
return (class_name, class_id, prob)
modules/pytorch_wrapper/pt_to_image.py
+2 -1
View File
@@ -6,7 +6,7 @@ class PtToImage:
"""
Casts a PyTorch tensor as an Image tensor.
category: PyTorch wrapper - Tensor data conversion
category: PyTorch wrapper - Image processing
"""
@classmethod
@@ -23,6 +23,7 @@ class PtToImage:
}
}
RETURN_NAMES: tuple = ("image",)
RETURN_TYPES: tuple = ("IMAGE",)
FUNCTION: str = "f"
CATEGORY: str = "Data Analysis"
modules/pytorch_wrapper/pt_to_image_transpose.py
+47
View File
@@ -0,0 +1,47 @@
from typing import Any, Dict
import torch
class PtToImageTranspose:
"""
Casts a PyTorch tensor as an Image tensor and transposes it from (C, H, W) to (H, W, C). For rank-4 inputs, the batch axis remains unchanged.
category: PyTorch wrapper - Image processing
"""
@classmethod
def INPUT_TYPES(cls) -> Dict[str, Any]:
"""
Defines the input types for the function.
Returns:
Dict[str, Any]: A dictionary specifying required input types.
"""
return {
"required": {
"tens": ("TENSOR", {}),
}
}
RETURN_NAMES: tuple = ("image",)
RETURN_TYPES: tuple = ("IMAGE",)
FUNCTION: str = "f"
CATEGORY: str = "Data Analysis"
def f(self, tens: torch.Tensor) -> tuple:
"""
Casts a PyTorch tensor as an Image tensor.
Args:
tens (torch.Tensor): PyTorch Tensor
Returns:
tuple: A tuple containing the image tensor.
"""
if tens.dim() == 3:
image = tens.permute(1, 2, 0)
elif tens.dim() == 4:
image = tens.permute(0, 2, 3, 1)
else:
raise ValueError("Only rank 3 or 4 tensors are supported.")
return (image,)
modules/pytorch_wrapper/pto_adam.py
+2 -2
View File
@@ -5,9 +5,9 @@ import torch.nn.functional as F
class PtoAdam:
"""
Converts PyTorch tensor to NumPy ndarray.
Instantiates the Adam optimizer.
category: PyTorch wrapper - Tensor data conversion
category: PyTorch wrapper - Training
"""
@classmethod
pyproject.toml
+1 -1
View File
@@ -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.10"
version = "2.0.11"
license = {file = "LICENSE"}
[project.urls]
tests/pytorch_wrapper/pt_crop_test.py
+49
View File
@@ -0,0 +1,49 @@
import os
import sys
import unittest
import torch
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.pt_crop import PtCrop
class TestPtCrop(unittest.TestCase):
def setUp(self):
"""Set up test instance."""
self.node = PtCrop()
# Test cases: (Input tensor, Height, Width, Expected shape)
self.test_cases = [
(torch.rand(1, 3, 16, 16), 8, 8, (1, 3, 8, 8)), # Crop 4D tensor
(torch.rand(3, 16, 16), 8, 8, (3, 8, 8)), # Crop 3D tensor
(torch.rand(1, 3, 16, 16), 16, 16, (1, 3, 16, 16)), # No crop, same size
]
def test_crop(self):
"""Test tensor cropping to specific sizes."""
for tens, height, width, expected_shape in self.test_cases:
with self.subTest(height=height, width=width):
result, = self.node.f(tens, height, width)
# Ensure return value is a tensor
self.assertTrue(isinstance(result, torch.Tensor), "Output is not a tensor")
# Ensure result shape matches expected shape
self.assertEqual(result.shape, expected_shape, f"Expected shape {expected_shape}, got {result.shape}")
def test_invalid_tensor_rank(self):
"""Test that an invalid tensor rank raises an error."""
tens = torch.rand(16, 16) # 2D tensor (invalid)
with self.assertRaises(ValueError, msg="Only rank 3 or 4 tensors are supported for crop."):
self.node.f(tens, 8, 8)
def test_crop_exceeds_dimensions(self):
"""Test that cropping beyond tensor dimensions raises an error."""
tens = torch.rand(1, 3, 16, 16)
with self.assertRaises(ValueError, msg="Specified crop size exceeds tensor dimensions."):
self.node.f(tens, 32, 32)
if __name__ == "__main__":
unittest.main()
tests/pytorch_wrapper/pt_from_image_transpose_test.py
+57
View File
@@ -0,0 +1,57 @@
import os
import sys
import unittest
import torch
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.pt_from_image_transpose import PtFromImageTranspose
class TestPtFromImageTranspose(unittest.TestCase):
def setUp(self):
"""Set up test instance."""
self.node = PtFromImageTranspose()
# Define test cases: (input_tensor, expected_shape)
self.test_cases = [
# Case 1: 3D tensor (H, W, C) -> (C, H, W)
(torch.ones(64, 128, 3), (3, 64, 128)), # (64,128,3) -> (3,64,128)
# Case 2: 4D tensor (N, H, W, C) -> (N, C, H, W)
(torch.ones(10, 64, 128, 3), (10, 3, 64, 128)), # (10,64,128,3) -> (10,3,64,128)
# Case 3: Invalid input (2D tensor)
(torch.ones(64, 128), ValueError), # 2D tensors should raise an error
# Case 4: Invalid input (5D tensor)
(torch.ones(10, 64, 128, 3, 2), ValueError), # 5D tensors should raise an error
]
def test_transpose(self):
"""Test image transposition operation with valid and invalid cases."""
for input_tensor, expected in self.test_cases:
with self.subTest(input_shape=input_tensor.shape):
if isinstance(expected, tuple):
# Valid cases
result, = self.node.f(input_tensor)
# Ensure return value is a tensor
self.assertTrue(isinstance(result, torch.Tensor), "Output is not a tensor")
# Ensure result has correct shape
self.assertEqual(result.shape, torch.Size(expected), f"Expected shape {expected}, got {result.shape}")
# Ensure result has the same dtype
self.assertEqual(result.dtype, input_tensor.dtype, "Dtype mismatch after transpose operation")
else:
# Invalid cases should raise an error
with self.assertRaises(expected):
self.node.f(input_tensor)
if __name__ == "__main__":
unittest.main()
tests/pytorch_wrapper/pt_interpolate_by_scale_factor_test.py
+60
View File
@@ -0,0 +1,60 @@
import os
import sys
import unittest
import torch
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.pt_interpolate_by_scale_factor import PtInterpolateByScaleFactor
class TestPtInterpolateByScaleFactor(unittest.TestCase):
def setUp(self):
"""Set up test instance."""
self.node = PtInterpolateByScaleFactor()
# Test cases 3d: (Input tensor, Scale factor, Mode, Expected shape)
self.test_cases_3d = [
(torch.rand(3, 16, 16), 2.0, "nearest", (3, 32, 32)),
(torch.rand(3, 16, 16), 2.0, "nearest-exact", (3, 32, 32)),
(torch.rand(3, 16, 16), 0.5, "bilinear", (3, 8, 8)),
(torch.rand(3, 16, 16), 1.5, "bicubic", (3, 24, 24)),
]
# Test cases: (Input tensor, Scale factor, Mode, Expected shape)
self.test_cases_4d = [
(torch.rand(1, 3, 16, 16), 2.0, "nearest", (1, 3, 32, 32)), # Upscale 4D tensor
(torch.rand(1, 3, 16, 16), 2.0, "nearest-exact", (1, 3, 32, 32)), # Upscale 4D tensor
(torch.rand(1, 3, 16, 16), 0.5, "bilinear", (1, 3, 8, 8)), # Downscale 4D tensor
(torch.rand(1, 3, 16, 16), 1.5, "bicubic", (1, 3, 24, 24)), # Bicubic mode with 4D tensor
]
def test_interpolation_3d(self):
"""Test tensor interpolation with different scale factors and modes."""
for tens, scale_factor, mode, expected_shape in self.test_cases_3d:
with self.subTest(scale_factor=scale_factor, mode=mode):
result, = self.node.f(tens, scale_factor, mode)
# Ensure return value is a tensor
self.assertTrue(isinstance(result, torch.Tensor), "Output is not a tensor")
# Ensure result shape matches expected shape
self.assertEqual(result.shape, expected_shape, f"Expected shape {expected_shape}, got {result.shape}")
def test_interpolation_4d(self):
"""Test tensor interpolation with different scale factors and modes."""
for tens, scale_factor, mode, expected_shape in self.test_cases_4d:
with self.subTest(scale_factor=scale_factor, mode=mode):
result, = self.node.f(tens, scale_factor, mode)
# Ensure return value is a tensor
self.assertTrue(isinstance(result, torch.Tensor), "Output is not a tensor")
# Ensure result shape matches expected shape
self.assertEqual(result.shape, expected_shape, f"Expected shape {expected_shape}, got {result.shape}")
if __name__ == "__main__":
unittest.main()
tests/pytorch_wrapper/pt_interpolate_to_size_test.py
+61
View File
@@ -0,0 +1,61 @@
import os
import sys
import unittest
import torch
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.pt_interpolate_to_size import PtInterpolateToSize
class TestPtInterpolateToSize(unittest.TestCase):
def setUp(self):
"""Set up test instance."""
self.node = PtInterpolateToSize()
# Test cases 3d: (Input tensor, Height, Width, Mode, Expected shape)
self.test_cases_3d = [
(torch.rand(3, 16, 16), 32, 32, "nearest", (3, 32, 32)), # Upscale 3D tensor
(torch.rand(3, 16, 16), 32, 32, "nearest-exact", (3, 32, 32)), # Upscale 3D tensor
(torch.rand(3, 16, 16), 8, 8, "bilinear", (3, 8, 8)), # Downscale 3D tensor
(torch.rand(3, 16, 16), 24, 24, "bicubic", (3, 24, 24)), # Bicubic mode with 3D tensor
]
# Test cases 4d: (Input tensor, Height, Width, Mode, Expected shape)
self.test_cases_4d = [
(torch.rand(1, 3, 16, 16), 32, 32, "nearest", (1, 3, 32, 32)), # Upscale 4D tensor
(torch.rand(1, 3, 16, 16), 32, 32, "nearest-exact", (1, 3, 32, 32)), # Upscale 4D tensor
(torch.rand(1, 3, 16, 16), 8, 8, "bilinear", (1, 3, 8, 8)), # Downscale 4D tensor
(torch.rand(1, 3, 16, 16), 24, 24, "bicubic", (1, 3, 24, 24)), # Bicubic mode with 4D tensor
]
def test_interpolation_3d(self):
"""Test tensor interpolation with different target sizes and modes."""
for tens, height, width, mode, expected_shape in self.test_cases_3d:
with self.subTest(height=height, width=width, mode=mode):
result, = self.node.f(tens, height, width, mode)
# Ensure return value is a tensor
self.assertTrue(isinstance(result, torch.Tensor), "Output is not a tensor")
# Ensure result shape matches expected shape
self.assertEqual(result.shape, expected_shape, f"Expected shape {expected_shape}, got {result.shape}")
def test_interpolation_4d(self):
"""Test tensor interpolation with different target sizes and modes."""
for tens, height, width, mode, expected_shape in self.test_cases_4d:
with self.subTest(height=height, width=width, mode=mode):
result, = self.node.f(tens, height, width, mode)
# Ensure return value is a tensor
self.assertTrue(isinstance(result, torch.Tensor), "Output is not a tensor")
# Ensure result shape matches expected shape
self.assertEqual(result.shape, expected_shape, f"Expected shape {expected_shape}, got {result.shape}")
if __name__ == "__main__":
unittest.main()
tests/pytorch_wrapper/pt_pad_test.py
+50
View File
@@ -0,0 +1,50 @@
import os
import sys
import unittest
import torch
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.pt_pad import PtPad
class TestPtPad(unittest.TestCase):
def setUp(self):
"""Set up test instance."""
self.node = PtPad()
# Test cases: (Input tensor, Height, Width, Expected shape)
self.test_cases = [
(torch.rand(1, 3, 16, 16), 32, 32, (1, 3, 32, 32)), # Pad 4D tensor
(torch.rand(3, 16, 16), 32, 32, (3, 32, 32)), # Pad 3D tensor
(torch.rand(1, 3, 16, 16), 16, 16, (1, 3, 16, 16)), # No padding, same size
]
def test_padding(self):
"""Test tensor padding to specific sizes."""
for tens, height, width, expected_shape in self.test_cases:
with self.subTest(height=height, width=width):
result, = self.node.f(tens, height, width)
# Ensure return value is a tensor
self.assertTrue(isinstance(result, torch.Tensor), "Output is not a tensor")
# Ensure result shape matches expected shape
self.assertEqual(result.shape, expected_shape, f"Expected shape {expected_shape}, got {result.shape}")
def test_invalid_tensor_rank(self):
"""Test that an invalid tensor rank raises an error."""
tens = torch.rand(16, 16) # 2D tensor (invalid)
with self.assertRaises(ValueError, msg="Only rank 3 or 4 tensors are supported for padding."):
self.node.f(tens, 32, 32)
def test_padding_smaller_than_tensor(self):
"""Test that padding to a smaller size raises an error."""
tens = torch.rand(1, 3, 16, 16)
with self.assertRaises(ValueError, msg="Specified dimensions must be greater than or equal to tensor dimensions."):
self.node.f(tens, 8, 8)
if __name__ == "__main__":
unittest.main()
tests/pytorch_wrapper/pt_predict_classification_conv_model_test.py
+68
View File
@@ -0,0 +1,68 @@
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_resize import PtvTransformsResize
from pytorch_wrapper.pt_data_loader import PtDataLoader
from pytorch_wrapper.pt_predict_classification_model import PtPredictClassificationModel
from pytorch_wrapper.pt_load_model import PtLoadModel
class TestPtPrediction(unittest.TestCase):
def setUp(self):
"""Set up test instance."""
# Transform
compose_node = PtvTransformsResize()
transform = compose_node.f(256, 256)[0]
# Dataset
self.node = PtvImageFolderDataset()
self.dataset = self.node.f("dog_and_cat/val",
transform=transform)[0]
# Model
self.model_node = PtnConvModel()
self.model = self.model_node.f(
input_dim="(3, 256, 256)",
penultimate_dim=0,
output_dim=2,
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]"
)[0]
# Load weights
epochs = 20
self.load_model = PtLoadModel()
self.loaded_model = self.load_model.f(self.model,f"dog_cat_{epochs}_epochs_conv.pt")
# Instantiates evaluation node
self.predict = PtPredictClassificationModel()
def test_prediction(self):
"""
Tests prediction.
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.
"""
inputs, y = next(iter(self.dataset))
class_id_to_name_map = {0:"cat", 1:"dog"}
class_name, class_id, prob = \
self.predict.f(self.model, inputs, class_id_to_name_map, use_gpu=True)
self.assertTrue(prob > 0.8)
self.assertEqual(class_id, 0)
self.assertEqual(class_name, "cat")
if __name__ == "__main__":
unittest.main()
tests/pytorch_wrapper/pt_to_image_transpose_test.py
+54
View File
@@ -0,0 +1,54 @@
import os
import sys
import unittest
import torch
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.pt_to_image_transpose import PtToImageTranspose
class TestPtToImageTranspose(unittest.TestCase):
def setUp(self):
"""Set up test instance."""
self.node = PtToImageTranspose()
# Define test cases: (input_tensor, expected_shape)
self.test_cases = [
# Case 1: 3D tensor (C, H, W) -> (H, W, C)
(torch.ones(3, 64, 128), (64, 128, 3)), # (3,64,128) -> (64,128,3)
# Case 2: 4D tensor (N, C, H, W) -> (N, H, W, C)
(torch.ones(10, 3, 64, 128), (10, 64, 128, 3)), # (10,3,64,128) -> (10,64,128,3)
# Case 3: Invalid input (2D tensor)
(torch.ones(64, 128), ValueError), # 2D tensors should raise an error
# Case 4: Invalid input (5D tensor)
(torch.ones(10, 3, 64, 128, 2), ValueError), # 5D tensors should raise an error
]
def test_transpose(self):
"""Test tensor to image transposition with valid and invalid cases."""
for input_tensor, expected in self.test_cases:
with self.subTest(input_shape=input_tensor.shape):
if isinstance(expected, tuple):
# Valid cases
result, = self.node.f(input_tensor)
# Ensure return value is a tensor
self.assertTrue(isinstance(result, torch.Tensor), "Output is not a tensor")
# Ensure result has correct shape
self.assertEqual(result.shape, torch.Size(expected), f"Expected shape {expected}, got {result.shape}")
# Ensure result has the same dtype
self.assertEqual(result.dtype, input_tensor.dtype, "Dtype mismatch after transpose operation")
else:
# Invalid cases should raise an error
with self.assertRaises(expected):
self.node.f(input_tensor)
if __name__ == "__main__":
unittest.main()