Function channel_first

code/dataset.py:92–95

(img)

    return img

def channel_first(img):
        if img.shape[-1] == 3:
            return rearrange(img, 'h w c -> c h w')
        return img