How do I classify colored images into 3 classes

[max22]

I'm struggling with the problem to write a program for classifying images (which are just colors) into 3 groups. I expect to obtain from each image a 3 vector [x,y,z].
The blue ones should be [1,0,0], green ones [0,1,0] and everything else beyond mixture of green and blue, such as red or yellow should yield [0,0,1].
The numbers are expected to be output of a softmax function and hence should be probabilities, each vector like this: [0.6,0.4,0]
for almost blue and with trace of green. I have written the following code so far but it gives a wrong results, see below, can anyone give me a hint on how to fix the code to yield proper vectors [x,y,1-x-y] ? There is apparently some bug as yellow gives this [0.0, 1.0, 0.0] instead of [0.1,0.2,0.7]. Thank you so much:

import tensorflow as tf
from tensorflow.keras import layers, models
from tensorflow.keras.preprocessing.image import ImageDataGenerator
import numpy as np
import os

# Define your image dimensions
img_width, img_height = 150, 150

# Define the path to your dataset
train_data_dir = 'C:/Users/Hynek/Desktop/HU/train'
test_data_dir = 'C:/Users/Hynek/Desktop/HU/test'  # Updated path for the test directory

# Build a CNN model with three classes
model = models.Sequential()
model.add(layers.Conv2D(32, (3, 3), activation='relu', input_shape=(img_width, img_height, 3)))
model.add(layers.MaxPooling2D((2, 2)))
model.add(layers.Conv2D(64, (3, 3), activation='relu'))
model.add(layers.MaxPooling2D((2, 2)))
model.add(layers.Conv2D(128, (3, 3), activation='relu'))
model.add(layers.MaxPooling2D((2, 2)))
model.add(layers.Flatten())
model.add(layers.Dense(128, activation='relu'))
model.add(layers.Dense(3, activation='softmax'))  # Change the output to 3 classes and use softmax activation

# Compile the model
model.compile(optimizer='adam', loss='sparse_categorical_crossentropy', metrics=['accuracy'])

# Use data augmentation for training
train_datagen = ImageDataGenerator(
    rescale=1./255,
    shear_range=0.2,
    zoom_range=0.2,
    horizontal_flip=True
)

train_generator = train_datagen.flow_from_directory(
    train_data_dir,
    target_size=(img_width, img_height),
    batch_size=32,
    class_mode='sparse',  # Use 'sparse' for integer-encoded labels
)

# Check if the generator has data
if len(train_generator) == 0:
    print("Error: No data found in the generator.")
else:
    # Train the model
    model.fit(train_generator, epochs=15, verbose=1, steps_per_epoch=len(train_generator))

    # Save the trained model weights
    model.save_weights('C:/Users/Hynek/Desktop/HU/trained_model_weights.h5')

    # Use data normalization for testing without data augmentation
    test_datagen = ImageDataGenerator(rescale=1./255)

    # Load test data without specifying classes
    test_generator = test_datagen.flow_from_directory(
        test_data_dir,
        target_size=(img_width, img_height),
        batch_size=32,
        class_mode=None,
        shuffle=False
    )

    # Check if the generator has data
    if len(test_generator) == 0:
        print("Error: No data found in the test generator.")
    else:
        # Make predictions on the test set
        predictions = model.predict(test_generator)

        # Define the class labels
        class_labels = ['1', '0', 'A']  # Adjusted class labels

        # Iterate through predictions and corresponding file names
        for prediction, file_name in zip(predictions, test_generator.filenames):
            # Reorder probabilities to match the adjusted class labels
            reordered_prediction = [prediction[1], prediction[0], prediction[2]]

            # Print the file name, the predicted class, and the probabilities for each class
            predicted_class = class_labels[np.argmax(reordered_prediction)]
            print(f"File: {file_name}, Predicted class: {predicted_class}, Probabilities: {reordered_prediction}")

Output:
File: 2\blue_image_1.jpg, Predicted class: 1, Probabilities: [1.0, 3.7408753e-32, 3.4386737e-35]
File: 2\border_color_image_1.jpg, Predicted class: 1, Probabilities: [1.0, 0.0, 0.0]
File: 2\border_color_image_10.jpg, Predicted class: 1, Probabilities: [1.0, 6.0347886e-11, 1.1029443e-32]
File: 2\border_color_image_11.jpg, Predicted class: 1, Probabilities: [0.98777896, 0.012221043, 1.5617029e-29]
File: 2\border_color_image_12.jpg, Predicted class: 0, Probabilities: [1.1639881e-06, 0.9999988, 1.8019694e-32]
File: 2\border_color_image_13.jpg, Predicted class: 0, Probabilities: [5.1332073e-15, 1.0, 3.781179e-38]
File: 2\border_color_image_14.jpg, Predicted class: 0, Probabilities: [6.26071e-22, 1.0, 0.0]
File: 2\border_color_image_15.jpg, Predicted class: 0, Probabilities: [1.1328224e-26, 1.0, 0.0]
File: 2\border_color_image_16.jpg, Predicted class: 0, Probabilities: [3.3592798e-31, 1.0, 0.0]
File: 2\border_color_image_17.jpg, Predicted class: 0, Probabilities: [8.1737104e-35, 1.0, 0.0]
File: 2\border_color_image_18.jpg, Predicted class: 0, Probabilities: [0.0, 1.0, 0.0]
File: 2\border_color_image_19.jpg, Predicted class: 0, Probabilities: [0.0, 1.0, 0.0]
File: 2\border_color_image_2.jpg, Predicted class: 1, Probabilities: [1.0, 0.0, 0.0]
File: 2\border_color_image_20.jpg, Predicted class: 0, Probabilities: [0.0, 1.0, 0.0]
File: 2\border_color_image_3.jpg, Predicted class: 1, Probabilities: [1.0, 0.0, 0.0]
File: 2\border_color_image_4.jpg, Predicted class: 1, Probabilities: [1.0, 0.0, 0.0]
File: 2\border_color_image_5.jpg, Predicted class: 1, Probabilities: [1.0, 0.0, 0.0]
File: 2\border_color_image_6.jpg, Predicted class: 1, Probabilities: [1.0, 3.5060138e-35, 0.0]
File: 2\border_color_image_7.jpg, Predicted class: 1, Probabilities: [1.0, 2.9402174e-30, 0.0]
File: 2\border_color_image_8.jpg, Predicted class: 1, Probabilities: [1.0, 1.0288415e-25, 0.0]
File: 2\border_color_image_9.jpg, Predicted class: 1, Probabilities: [1.0, 1.00328885e-17, 1.8057328e-35]
File: 2\green_image_2.jpg, Predicted class: 0, Probabilities: [0.0, 1.0, 0.0]
File: 2\red_image.jpg, Predicted class: 1, Probabilities: [0.9999999, 6.662962e-08, 3.89023e-20]
File: 2\yellow_image.jpg, Predicted class: 0, Probabilities: [0.0, 1.0, 0.0]