[Python] How to convert discrete choices (of a color) from neural network output to colors...

Context: I am learning to use tensorflow and want to do a simple experiment where I provide a neural network with 4 color choices for each pixel. The network should learn to pick the best colors from choices available for each pixel to represent the image it is given. Basically, a dithering neural network.

I need to map networks output, which comes in the shape of (width*height, 4) to an RGB image. That will be compared to blurred reference image.

I have the following code which I think does what I want, but over numpy matrices:

def convert_to_image(output_matrix, color_choices):
"""Maps network outputs to RGB colors using predefined choices."""
# Converts (height*width, 4-floats) to (height*width) array of 0-3 indices
chosen_colors = np.argmax(output_matrix, axis=1)
# Uses color_choices to mach each indice to a pixel color
mapped_image = color_choices[np.arange(color_choices.shape[0]), chosen_colors] * 255

mapped_image = mapped_image.reshape((INPUT_SHAPE[0], INPUT_SHAPE[1], 3)) # Reshape to (height, width, 3)
return mapped_image


However I soon learned I cannot use this for reinforcement learning, because tensorflow will not be able to calculate what it should change in the network.

So here is my tensorflow attempt and its output:

def convert_to_image_tf(output_matrices, color_choices_batch):
"""Maps network outputs to RGB colors using predefined choices in TensorFlow."""
batch_size = output_matrices.shape[0]
print("output_matrices dimensions: " + str(output_matrices.shape))
print("color_choices_batch dimensions: " + str(color_choices_batch.shape))
chosen_colors = tf.argmax(output_matrices, axis=2, output_type=tf.int32)
print("chosen_colors shape: " + str(chosen_colors.shape))
#batch_indices = tf.range(batch_size)[:, tf.newaxis] # Shape: (batch_size, 1)
#print("batch_indices shape: " + str(chosen_colors.shape))
#chosen_colors_expanded = tf.stack([batch_indices, chosen_colors], axis=-1) # Shape: (batch_size, width*height, 2)
#print("chosen_colors_expanded shape: " + str(chosen_colors_expanded.shape))

batch_size, num_pixels, num_choices_per_px, num_colors = color_choices_batch.shape # Extract dimensions

# Use batch-wise advanced indexing
mapped_image = color_choices[np.arange(batch_size)[:, None], np.arange(num_pixels), chosen_colors] * 255

#mapped_image = tf.gather_nd(color_choices_batch, chosen_colors) * 255 # Shape: (batch_size, width*height, 3)
# mapped_image = tf.gather(color_choices_batch, chosen_colors, batch_dims=0) * 255
print("mapped_image shape: " + str(mapped_image.shape))
mapped_image = tf.reshape(mapped_image, (INPUT_SHAPE[0], INPUT_SHAPE[1], 3)) # Reshape to (height, width, 3)


Output:

output_matrices dimensions: (32, 1024, 4)
color_choices_batch dimensions: (32, 1024, 4, 3)
chosen_colors shape: (32, 1024)
Traceback (most recent call last):
File "xxx\python_tests\dither_test_1.py", line 229, in <module>
loss = train_step(color_decider, image_batch, color_choices_batch)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "xxx\python_tests\dither_test_1.py", line 209, in train_step
generated_images = convert_to_image_tf(raw_outputs, color_choices_batch)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "xxx\python_tests\dither_test_1.py", line 182, in convert_to_image_tf
mapped_image = color_choices[np.arange(batch_size)[:, None], np.arange(num_pixels), chosen_colors] * 255
~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
IndexError: index 4 is out of bounds for axis 1 with size 4


Full program is below. Note that there seem to be some recent issue with tensorflow imports, so you might need to change those:

import tensorflow as tf
import tensorflow.python.keras as keras
import tensorflow.python.keras.layers as layers
import tensorflow.python.keras.losses as losses
import tensorflow.python.keras.optimizers as optimizers
import numpy as np
import os
import random
import cv2

# Define input shape
INPUT_SHAPE = (32, 32, 3) # Example: 32x32 RGB images
NUM_COLORS = 4 # The four color choices per pixel


# Second network: Color decision maker
def build_color_decider():
# First input: Matrix with 4 values per pixel
color_choices = keras.Input(shape=(INPUT_SHAPE[0] * INPUT_SHAPE[1], 4, 3), name="ColorChoices")

# Second input: Normal RGB image
image_data = keras.Input(shape=(INPUT_SHAPE[0], INPUT_SHAPE[1], 3), name="ImageData")

# Process color choices matrix
x1 = layers.Conv2D(32, (3,3), activation="relu")(color_choices)
x1 = layers.Flatten()(x1)

# Process image data matrix
x2 = layers.Conv2D(32, (3,3), activation="relu")(image_data)
x2 = layers.Flatten()(x2)

merged = layers.Concatenate()([x1, x2])
merged = layers.Dense(128, activation="relu")(merged)
# One number per pixel, representing the chosen color
output = layers.Dense(INPUT_SHAPE[0] * INPUT_SHAPE[1]*4, activation='linear')(merged) # Example output layer
output = layers.Reshape((INPUT_SHAPE[0] * INPUT_SHAPE[1], 4))(output)
model = keras.Model(inputs=[color_choices, image_data], outputs=output)
model.compile(optimizer="adam", loss="categorical_crossentropy", metrics=["accuracy"])

model.summary()
# model = keras.Sequential([
# layers.Input(shape=(INPUT_SHAPE[0], INPUT_SHAPE[1], 4)),
# layers.Conv2D(INPUT_SHAPE[0], (3,3), activation='relu', input_shape=INPUT_SHAPE),
# layers.Conv2D(64, (3,3), activation='relu'),
# layers.Flatten(),
# layers.Dense(1024, activation='relu'),
# layers.Dense(INPUT_SHAPE[0] * INPUT_SHAPE[1] * 3, activation='linear') # Final RGB values
# ])
return model

image_dirs = ["ENTER IMAGE FOLDER HERE"] # Example folders

def load_random_image():
"""Loads a random image from one of the folders."""
chosen_dir = random.choice(image_dirs) # Pick a random folder
images = os.listdir(chosen_dir) # List images
img_path = os.path.join(chosen_dir, random.choice(images)) # Pick a random image
return img_path

def preprocess_image(img_path, downsample_factor=2, target_size=(32, 32)):
"""Loads, downsamples, and crops the image."""
img = cv2.imread(img_path) # Read image
# Calculate new size so that smallest dimension is same as target size
h, w, _ = img.shape
if h < w:
new_h = target_size[0]
new_w = int(w * (new_h / h))
else:
new_w = target_size[0]
new_h = int(h * (new_w / w))
img = cv2.resize(img, (new_w, new_h)) # Resize to maintain aspect ratio
#img = cv2.resize(img, (64, 64)) # Downsample (adjust size)

# Random Crop (Center Crop for simplicity)
crop_size = target_size[0] # Assume square crop
h, w, _ = img.shape
start_x = (w - crop_size) // 2
start_y = (h - crop_size) // 2
img = img[start_y:start_y+crop_size, start_x:start_x+crop_size] # Crop

img = img / 255.0 # Normalize
return img

def data_generator(batch_size=32):
"""Yields batches of processed images."""
while True:
batch_images = []
for _ in range(batch_size):
img_path = load_random_image()
img = preprocess_image(img_path)
batch_images.append(img)

yield np.array(batch_images)

# Create the models
color_decider = build_color_decider()

color_map = {
"cyan": [0, 255, 255], # RGB for Cyan
"magenta": [255, 0, 255], # RGB for Magenta
"yellow": [255, 255, 0], # RGB for Yellow
"black": [0, 0, 0], # RGB for Black
"white": [255, 255, 255], # RGB for White
"green": [0, 255, 0], # RGB for Green
"red": [255, 0, 0], # RGB for Red
"blue": [0, 0, 255] # RGB for Blue
}

# normalize color map values to [0, 1] range
color_map = {k: np.array(v) / 255.0 for k, v in color_map.items()}

pixel_option_1 = np.array([
color_map["cyan"],
color_map["magenta"],
color_map["yellow"],
color_map["black"]
])

pixel_option_2 = np.array([
color_map["blue"],
color_map["black"],
color_map["red"],
color_map["white"]
])

# first run with preset color choices - CMYK
color_choices = np.array([
pixel_option_1 if xcolor%2==0 else pixel_option_2 for xcolor in range(INPUT_SHAPE[0]*INPUT_SHAPE[1])
])

def data_generator_with_color_choices(batch_size=32):
"""Yields batches of processed images with color choices."""

images_generator = data_generator(batch_size)
while True:
batch_color_choices = [color_choices for _ in range(batch_size)]

yield np.array(batch_color_choices), np.array(next(images_generator))


def convert_to_image_tf(output_matrices, color_choices_batch):
"""Maps network outputs to RGB colors using predefined choices in TensorFlow."""
batch_size = output_matrices.shape[0]
print("output_matrices dimensions: " + str(output_matrices.shape))
print("color_choices_batch dimensions: " + str(color_choices_batch.shape))
chosen_colors = tf.argmax(output_matrices, axis=2, output_type=tf.int32)
print("chosen_colors shape: " + str(chosen_colors.shape))
#batch_indices = tf.range(batch_size)[:, tf.newaxis] # Shape: (batch_size, 1)
#print("batch_indices shape: " + str(chosen_colors.shape))
#chosen_colors_expanded = tf.stack([batch_indices, chosen_colors], axis=-1) # Shape: (batch_size, width*height, 2)
#print("chosen_colors_expanded shape: " + str(chosen_colors_expanded.shape))

batch_size, num_pixels, num_choices_per_px, num_colors = color_choices_batch.shape # Extract dimensions

# Use batch-wise advanced indexing
mapped_image = color_choices[np.arange(batch_size)[:, None], np.arange(num_pixels), chosen_colors] * 255

#mapped_image = tf.gather_nd(color_choices_batch, chosen_colors) * 255 # Shape: (batch_size, width*height, 3)
# mapped_image = tf.gather(color_choices_batch, chosen_colors, batch_dims=0) * 255
print("mapped_image shape: " + str(mapped_image.shape))
mapped_image = tf.reshape(mapped_image, (INPUT_SHAPE[0], INPUT_SHAPE[1], 3)) # Reshape to (height, width, 3)

return mapped_image

def blur_image(image):
"""Applies Gaussian blur to smooth the image."""
return cv2.GaussianBlur(image, (5, 5), 0) # Kernel size of (5,5)

train_data = tf.data.Dataset.from_generator(
lambda: data_generator_with_color_choices(32),
output_types=(tf.float32, tf.float32), # Two inputs: color choices & image data
output_shapes=((None, INPUT_SHAPE[0]*INPUT_SHAPE[1], 4, 3), (None, INPUT_SHAPE[0], INPUT_SHAPE[1], 3)) # First is color choices, second is image data
)

loss_fn = tf.keras.losses.MeanSquaredError()
optimizer = tf.keras.optimizers.Adam(learning_rate=0.001)

# Custom training loop
def train_step(model, image_batch, color_choices_batch):
with tf.GradientTape() as tape:
batch_size = image_batch.shape[0]
raw_outputs = model((color_choices_batch, image_batch)) # Get predictions
generated_images = convert_to_image_tf(raw_outputs, color_choices_batch)

blurred_generated_images = tf.nn.avg_pool2d(generated_images, ksize=5, strides=1, padding="SAME")
blurred_original_images = tf.nn.avg_pool2d(image_batch, ksize=5, strides=1, padding="SAME")

# Compute loss
loss = loss_fn(blurred_original_images, blurred_generated_images)

# Compute and apply gradients
grads = tape.gradient(loss, model.trainable_variables)
optimizer.apply_gradients(zip(grads, model.trainable_variables))

return loss.numpy()

# Training loop
for epoch in range(300):
for step, (color_choices_batch, image_batch) in enumerate(train_data.take(100)):
loss = train_step(color_decider, image_batch, color_choices_batch)
print(f"Epoch {epoch}, Step {step}, Loss: {loss}")

Continue reading...