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]