31 lines
1.1 KiB
Python
31 lines
1.1 KiB
Python
import gzip
|
|
import numpy as np
|
|
|
|
def read_images(filename):
|
|
with gzip.open(filename, 'rb') as f:
|
|
f.read(16)
|
|
return np.frombuffer(f.read(), dtype=np.uint8).reshape(-1, 28, 28)
|
|
|
|
def pixel_to_stimulation(pixel_value, max_stimulation=5.0):
|
|
return int(pixel_value / 255.0 * max_stimulation * 51) # Convert to 0-255 range
|
|
|
|
def read_labels(filename):
|
|
with gzip.open(filename, 'rb') as f:
|
|
f.read(8)
|
|
return np.frombuffer(f.read(), dtype=np.uint8)
|
|
|
|
train_images = read_images('dataset/train-images-idx3-ubyte.gz')
|
|
train_labels = read_labels('dataset/train-labels-idx1-ubyte.gz')
|
|
|
|
indices_of_nines = np.where(train_labels == 9)[0]
|
|
image_of_nine = train_images[indices_of_nines[0]]
|
|
|
|
stimulation_pattern = np.vectorize(pixel_to_stimulation)(image_of_nine)
|
|
|
|
# Output the stimulation pattern as a compact byte array
|
|
with open('stimulation_pattern.h', 'w') as f:
|
|
f.write("const PROGMEM uint8_t stimulation_pattern[] = {")
|
|
f.write(", ".join(map(str, stimulation_pattern.flatten())))
|
|
f.write("};\n")
|
|
f.write(f"const uint16_t PATTERN_SIZE = {stimulation_pattern.size};\n")
|