There are a few helper functions we need inside of class—we're going to go over the steps to create them and why we need them here:

1. The array_to_color function is a helper function to the translate function in step 2—it essentially simplifies one of our steps to map our points to a unique color mapping:

# Translate data to color
def array_to_color(self,array, cmap="Oranges"):
    s_m = plt.cm.ScalarMappable(cmap=cmap)
return s_m.to_rgba(array)[:,:-1]

2. The translate function provides functionality to convert the points in the 3D MNIST data into a pixel x, y, z ,with a three-color definition associated with it:

def translate(self,x):
    xx = np.ndarray((x.shape[0], 4096, 3))
for i in range(x.shape[0]):
    xx[i] = self.array_to_color(x[i])
del x
return xx

3. The load_2d_encoded_MNIST function reads in the npy files we created from our encoder in the Encoding 2D data recipe:

def load_2D_encoded_MNIST(self):
    (_, self.Y_train_2D), (_, self.Y_test_2D) = mnist.load_data()
    self.X_train_2D_encoded = np.load('x_train_encoded.npy')
    self.X_test_2D_encoded = np.load('x_test_encoded.npy')
return

4. The next method, load_3D_MNIST, takes the input of an input dir and loads the 3D MNIST data—the first step is to grab the raw data from the h5 file:

def load_3D_MNIST(self,input_dir):
     raw = h5py.File(input_dir, 'r')

5. We're going to grab the X_train data, normalize the data, and reshape it to our input that we need for our models:

self.X_train_3D = np.array(raw['X_train'])
self.X_train_3D = ( np.float32(self.X_train_3D) - 127.5) / 127.5
self.X_train_3D = self.translate(self.X_train_3D).reshape(-1, 16, 16, 16, 3)

6. We'll repeat the same process for the test set—import the test set data, normalize, and then reshape to match the model input:

self.X_test_3D = np.array(raw['X_test'])
self.X_test_3D = ( np.float32(self.X_test_3D) - 127.5) / 127.5
self.X_test_3D = self.translate(self.X_test_3D).reshape(-1, 16, 16, 16, 3)

7. We'll save the Y datasets as they are not matched to their 2D counterparts:

self.Y_train_3D = np.array(raw['y_train'])
self.Y_test_3D = np.array(raw['y_test'])
return

Let's move on to the training method!