CSV 3D data matrix to normalized pixel matrix

csv_matrix_to_image.py

#!/usr/bin/env python3
import csv
import numpy as np
import matplotlib as mpl
from PIL import Image

MATRIX_SIZE = 64
NORMALIZE_BIAS = 255 # 0.0 - 1.0
Z_DEPTH = 3

"""
Expects data like:
x y z
1 1 4536.00
1 2 90.00
1 3 88.00
1 4 94.00
1 5 94.00
To use floats throughout, modify the x, y, z line
"""

if __name__ == '__main__':
    input_file  = "./input/heatmap_data.txt"
    # 64x64xZ_DEPTH matrix
    data_matrix = [[[0] * Z_DEPTH for _ in range(MATRIX_SIZE)] for _ in range(MATRIX_SIZE)]
    with open(input_file) as csvfile:
        reader = csv.DictReader(csvfile, delimiter=" ")
        for row in reader:
            x, y, z  = map(int, map(float, row.values()))
            r, g, b, a = mpl.colormaps["viridis"](np.float32(z) / NORMALIZE_BIAS)
            data_matrix[x-1][y-1][:Z_DEPTH] = r, g, b  # type: ignore
    # copy python matrix to numpy flat array
    data_matrix_array = np.array(data_matrix, dtype=np.float32, copy=True)
    data_matrix_nd_array = np.ndarray([MATRIX_SIZE, MATRIX_SIZE, Z_DEPTH], dtype=np.float32, buffer=data_matrix_array)
    image = Image.fromarray(np.uint8(data_matrix_array * NORMALIZE_BIAS))
    image.convert('RGB')
    image.show()