mantel/create_image.py

import numpy as np
import matplotlib.pyplot as plt

from models import train_kohonen_network, create_kohonen_params

X = np.random.random((1000,3))

IMAGE_SIZE = 100

params = create_kohonen_params(X, IMAGE_SIZE, IMAGE_SIZE, 100).unwrap()
image_data = train_kohonen_network(X, params, use_mlflow=True).unwrap()

plt.imsave(f'{IMAGE_SIZE}.png', image_data)