import json


def pre_proc(batch, **kwargs):
    """
    Pre-processing of batch.
    Assumes format (n_samples, n_channels, height, width).
    """
    #opening stats file
    with open(kwargs["stats_fp"]) as f:
        stats = json.load(f)

    #converting keys from str to int
    stats = {int(k): v for k, v in stats.items()}

    #normalizing per channel
    for k, s in stats.items():
        batch[:, k, :, :] = (batch[:, k, :, :] - s["mean"])/s["std"]

    return batch