Nodes/ComfyUI-Pt-Wrapper/Ptv Transforms Data Augment
ComfyUI Node

Ptv Transforms Data Augment

Make your tiny image dataset last longer

By HowToSD·Created about a year ago·Updated about a year ago· 7
Ptv Transforms Data Augment
    • PTVTRANSFORM
    h_flip_prob0.000
    v_flip_prob0.000
    rotate_degree0.000
    h_translate_ratio0.000
    v_translate_ratio0.000
    min_scale1.000
    max_scale1.000

    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:

    • RandomHorizontalFlip and/or RandomVerticalFlip, added only if the corresponding probability is above zero
    • a RandomAffine with 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.5 is 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.1 is a small shift.
    • min_scale / max_scale - zoom range, e.g. 0.91.1. If both are 1.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 the ToTensor step. Keep the source a plain ImageFolder.
    • 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.0 and both scales at 1.0 builds a pipeline that only does ToTensor. 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.

    CategoryTraining

    Inputs (7)

    NameTypeDefaultDescription
    h_flip_probFLOAT0.0000–1
    v_flip_probFLOAT0.0000–1
    rotate_degreeFLOAT0.000-360–360
    h_translate_ratioFLOAT0.0000–1
    v_translate_ratioFLOAT0.0000–1
    min_scaleFLOAT1.0000.5–1
    max_scaleFLOAT1.0001–2

    Outputs (1)

    NameTypeDescription
    PTVTRANSFORMPTVTRANSFORM