test_aug.py

import tensorpack
from tensorpack.dataflow import (FakeData, MultiThreadMapData, MultiProcessMapData, imgaug, BatchData)
import numpy as np
import time

def test(x):
    # simulate heavy workload
    time.sleep(0.01)
    
    x[0], _prms = aug.augment_return_params(x[0])
    print (_prms)
    
    return x

ds = FakeData([[10, 10, 3]])
aug = imgaug.Gamma()
ds = MultiProcessMapData(ds, 4, test)
ds = BatchData(ds, 4)
ds.reset_state()
for i, dp in enumerate(ds):
    if (i+1) % 2 == 0:
        exit(0)
        
# output:
# 
# -0.43118484206945873
# -0.43118484206945873
# -0.43118484206945873
# -0.43118484206945873
# -0.45425744974768434
# -0.45425744974768434
# -0.45425744974768434
# -0.45425744974768434