augment.py

# -*- Python -*-

# Perform imagenet-style augmentation and normalization on the shards 
# on stdin, returning a new dataset on stdout.

import sys
from torchvision import transforms
import webdataset as wds

normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])


augment = transforms.Compose(
    [
        transforms.RandomResizedCrop(224),
        transforms.RandomHorizontalFlip(),
        transforms.ToTensor(),
        normalize,
    ]
)

dataset = wds.WebDataset("-").decode("pil")

sink = wds.TarWriter("-")
for sample in dataset:
    print(sample.get("__key__"), file=sys.stderr)
    sample["npy"] = augment(sample["jpg"]).numpy().astype("float16")
    sink.write(sample)