sample swap for image (batch, cannel, height, width*2)

sample_swap.py

def sample_swap(shape):
    b, c, h, w = shape
    left = np.arange(0, w//2)
    right = np.arange(w//2, w)
    flip = np.random.uniform(size=(b, 1)) > .5
    b = np.arange(b)[:, None, None, None]
    c = np.arange(c)[None, :, None, None]
    h = np.arange(h)[None, None, :, None]
    w = np.concatenate([left+flip*(w//2), right-flip*(w//2)], -1)
    w = w[:, None, None, :]
    return np.broadcast_arrays(b, c, h, w)