Images get distorted after splitting the array using scikit-learn

Question

I have a .mat file that contains an array of 600 24*21 images. Hence the resulting dimension is (24, 21, 600). The follow