fastnoise.py

import random
from random import choice

import tensorflow as tf
import numpy as np
import pyfastnoisesimd as fns

def ch3_fractal_noise(size_chw, seed,
    color_gain_min=0.5, color_gain_max=1.0,
    octaves_min=3, octaves_max=6,
    frequency_min=0.01, frequency_max=0.075,
    perturb_frequency_min=0.2, perturb_frequency_max=0.7,
    fractal_gain_min=0.4, fractal_gain_max=0.6,
    ):

    c, h, w = size_chw

    noise_g = fns.Noise(seed=seed, numWorkers=fns.num_virtual_cores())

    # basic params
    noise_type = choice(
        [e for e in fns.NoiseType if e != fns.NoiseType.WhiteNoise])
    frequency = np.random.uniform(frequency_min, frequency_max)

    noise_g.noiseType = noise_type
    noise_g.frequency = frequency

    # fractal params
    octaves = choice(range(octaves_min, octaves_max+1))
    fractal_gain = np.random.uniform(fractal_gain_min, fractal_gain_max)
    fractal_type = choice([e for e in fns.FractalType])

    noise_g.fractal.octaves = octaves 
    noise_g.fractal.lacunarity = 2.1
    noise_g.fractal.gain = fractal_gain
    noise_g.fractal.type = fractal_type

    # cellular params
    cellular_distance_func = choice(
        [e for e in fns.CellularDistanceFunction])
    cellular_lookup_frequence = np.random.uniform(0.1, 0.3)
    cellular_jittar = np.random.uniform(0.45, 0.9)
    cellular_return_type = choice(
        [e for e in fns.CellularReturnType])

    noise_g.cellularDistanceFunc = cellular_distance_func
    noise_g.cellularJittar = cellular_jittar
    noise_g.cellularLookupFrequence = cellular_lookup_frequence
    noise_g.cellularReturnType = cellular_return_type

    # perturbation params
    perturb_type = choice(
        [fns.PerturbType.NoPerturb, fns.PerturbType.Gradient, fns.PerturbType.GradientFractal])
    perturb_amp = np.random.uniform(0.6, 1.0)
    perturb_octaves = choice(range(octaves_min, octaves_max+1))
    perturb_frequency = np.random.uniform(perturb_frequency_min, perturb_frequency_max)
    perturb_gain = np.random.uniform(0.5, 0.55)

    noise_g.perturb.perturbType = perturb_type
    noise_g.perturb.frequency = perturb_frequency
    noise_g.perturb.octaves = perturb_octaves
    noise_g.perturb.amp = perturb_amp
    noise_g.perturb.gain = perturb_gain
    noise_g.perturb.normaliseLength = np.random.uniform(0.5, 1.5)
    
    # generate noise image
    noise_tmp = noise_g.genAsGrid([1,h*c,w])
    noise = np.zeros([c,h,w])
    for ch in range(c):
        noise[ch] = noise_tmp[0,h*ch:h*(ch+1),:]

    # color adjustment
    noise = (noise / 2) + 0.5
    noise_amax = np.amax(noise, axis=(1,2))
    noise /= (noise_amax[:, np.newaxis, np.newaxis] + 1e-12)

    color_gain = np.random.uniform(color_gain_min, color_gain_max, size=(c,))
    noise *= color_gain[:, np.newaxis, np.newaxis]

    return noise


def create_fractal_noise_dataset(size_hwc, seed=0, num_parallel_calls=8):
    h,w,c = size_hwc
    size_chw = (c,h,w)

    def _get_generator(seed_v):
        def _generator():
            _seed = seed_v
            random.seed(seed_v + 5)
            np.random.seed(seed_v // 256 + 11)
            while True:
                yield ch3_fractal_noise(size_chw, _seed)
                _seed += 1
        return _generator

    def _batch(*x):
        c, h, w = x[0].shape
        n = len(x)
        x = tf.concat(x, axis=0)
        x = tf.reshape(x, (n, c, h, w))
        return x

    dataset = tuple(
        tf.data.Dataset.from_generator(
            _get_generator((i+1)*(2**28)+seed),
            output_signature=tf.TensorSpec(shape=size_chw, dtype=tf.float32))
        for i in range(num_parallel_calls))
    dataset = tf.data.Dataset.zip(dataset)
    dataset = dataset.map(_batch)
    dataset = dataset.unbatch()
    return dataset

def main():
    from PIL import Image
    dataset = create_fractal_noise_dataset([1024,512,3])
    dataset = dataset.map(lambda x: tf.transpose(x, (1,2,0)))
    dataset = dataset.batch(16)
    dataset = dataset.prefetch(1)

    for i, d in zip(range(64), dataset):
        d_np = d.numpy()
        d_np = (d_np * 255).astype(np.uint8)
        for j in range(16):
            image = Image.fromarray(d_np[j])
            image.save(f'out_{i:04d}_{j:02d}.jpg')


if __name__ == '__main__':
    main()