Ptv Transforms Data Augment
Make your tiny image dataset last longer
- PTVTRANSFORM
Small datasets overfit. That's the problem this node exists to fight. When you're training a classifier on a couple of hundred photos of your own, the model will happily memorize them instead of learning what a dog is. Data augmentation fakes a bigger dataset by warping each image slightly - flipping it, rotating it, shifting it - so the model sees "the same" image from a dozen angles and learns the underlying pattern instead of the pixel noise.
How it works
This node doesn't touch any data itself. It builds a torchvision.transforms.Compose pipeline - a PTVTRANSFORM - that you then plug into the transform input of Ptv Image Folder Dataset. Under the hood it assembles:
RandomHorizontalFlipand/orRandomVerticalFlip, added only if the corresponding probability is above zero- a
RandomAffinewith your rotation, translation and scale settings - a final
ToTensor()to convert images to tensors in[0, 1]
Because the flips are conditional on a probability, leaving a flip at 0.0 simply omits it from the pipeline. Each epoch, every image gets a fresh random warp - that's the augmentation.
The inputs that matter
All seven are floats with sensible defaults of "off", so start small:
- h_flip_prob / v_flip_prob - chance (0–1) of a horizontal/vertical flip.
0.5is the classic start for horizontal. - rotate_degree - maximum rotation in degrees, e.g.
15. Negative allowed, but magnitude is what matters. - h_translate_ratio / v_translate_ratio - max translation as a fraction of image size.
0.1is a small shift. - min_scale / max_scale - zoom range, e.g.
0.9–1.1. If both are1.0, no scaling is applied.
Output is a single PTVTRANSFORM - wire it into Ptv Image Folder Dataset's transform input.
The traps
Don't flip everything, and don't flip down. Horizontal flips are safe for most real-world photos. Vertical flips produce nonsense for most scenes (trees upside down) and can actively hurt - only use v_flip_prob when your domain is rotation-invariant, like certain medical or satellite imagery. RandomAffine rotation beyond ~15-20° also starts degrading real photos. More aggressive is not automatically better.
Augmentation is for training, not evaluation. Wire the plain Ptv Transforms To Tensor (or resize) into your evaluation dataset, and the augment pipeline only into training. If you augment your validation set, your accuracy numbers mean nothing.
Order of operations: if you're also resizing, do it via a separate resize transform in the chain - the augment node applies ToTensor itself, so don't stack a second one.
Install
Standard pack install:
cd ComfyUI/custom_nodes
git clone https://github.com/HowToSD/ComfyUI-Pt-Wrapper
Restart ComfyUI, or find "ComfyUI-Pt-Wrapper" in ComfyUI Manager. No extra dependencies.
Common issues
- "Expected a PIL Image" errors - the pipeline expects PIL inputs from
ImageFolder; feeding it already-tensorized data breaks theToTensorstep. Keep the source a plainImageFolder. - Validation accuracy collapsed - you probably augmented your eval set, or flipped vertical on a dataset that can't take it.
- Nothing seems to change - all probabilities at
0.0and both scales at1.0builds a pipeline that only doesToTensor. That's correct behavior, not a bug.
Start with h_flip_prob: 0.5 and a rotate_degree: 10, run a training pass, and watch the validation curve for signs of overfitting before you crank it further.
Inputs (7)
| Name | Type | Default | Description |
|---|---|---|---|
| h_flip_prob | FLOAT | 0.0000–1 | — |
| v_flip_prob | FLOAT | 0.0000–1 | — |
| rotate_degree | FLOAT | 0.000-360–360 | — |
| h_translate_ratio | FLOAT | 0.0000–1 | — |
| v_translate_ratio | FLOAT | 0.0000–1 | — |
| min_scale | FLOAT | 1.0000.5–1 | — |
| max_scale | FLOAT | 1.0001–2 | — |
Outputs (1)
| Name | Type | Description |
|---|---|---|
| PTVTRANSFORM | PTVTRANSFORM | — |