diff --git a/Network.py b/Network.py index d82394b..1eb259b 100644 --- a/Network.py +++ b/Network.py @@ -9,9 +9,12 @@ import PIL import PIL.Image import numpy + +def normalize(image, label): + return image / 255, label + # Set the paths to the folders containing the training data train_data_dir = "Training/" -validation_data_dir = "Validation/" # Set the number of classes and batch size @@ -20,7 +23,7 @@ batch_size = 32 # Set the image size and input shape img_width, img_height = 100, 100 -input_shape = (img_width, img_height, 3) +input_shape = (img_width, img_height, 1) train_ds = tf.keras.utils.image_dataset_from_directory( train_data_dir, @@ -42,3 +45,36 @@ val_ds = tf.keras.utils.image_dataset_from_directory( class_names = train_ds.class_names print(class_names) + +train_ds = train_ds.map(normalize) +val_ds = val_ds.map(normalize) + +# Define the model architecture +model = tf.keras.Sequential([ + layers.Conv2D(16, 3, padding='same', activation='relu', input_shape=(img_height, img_width, 1)), + layers.MaxPooling2D(), + layers.Conv2D(32, 3, padding='same', activation='relu'), + layers.MaxPooling2D(), + layers.Conv2D(64, 3, padding='same', activation='relu'), + layers.MaxPooling2D(), + layers.Flatten(), + layers.Dense(128, activation='relu'), + layers.Dense(num_classes, activation='softmax') +]) + +# Compile the model +model.compile(optimizer='adam', + loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True), + metrics=['accuracy']) + +model.summary() + + +epochs=10 +# Train the model +model.fit(train_ds, + validation_data=val_ds, + epochs=epochs) + +# Save the trained model +model.save('trained_model.h5') diff --git a/Tester2.py b/Tester2.py new file mode 100644 index 0000000..f6dd52e --- /dev/null +++ b/Tester2.py @@ -0,0 +1,38 @@ +import os +import numpy as np +import tensorflow as tf +from tensorflow import keras + +train_data_dir = "Training/" + +train_ds = tf.keras.utils.image_dataset_from_directory(train_data_dir, validation_split=0.2, + subset="training", seed=123, batch_size=32, image_size=(100, 100)) +val_ds = tf.keras.utils.image_dataset_from_directory(train_data_dir, validation_split=0.2, + subset="validation", seed=123, batch_size=32, image_size=(100, 100)) + +model = keras.models.load_model("trained_model") +predictions = model.predict(val_ds.take(32)) + + +classNames = ['Empty', 'Food','People'] + +# Make predictions +direct = '' +i = 0 +for image, _ in val_ds.take(32): + predicted_class_index = np.argmax(predictions[i]) + predicted_class = classNames[predicted_class_index] + filename = predicted_class + str(i) + '.jpeg' + tf.keras.preprocessing.image.save_img(direct+filename, image[0]) + print('Predicted class:', predicted_class) + i += 1 + +#direct = '' + +#i = 0 +#for image, _ in val_ds.take(32): +# predictedLabel = int(predictions[i] >= 0.5) +# +# filename = classNames[predictedLabel] + str(i) + '.jpeg' +# tf.keras.preprocessing.image.save_img(direct+filename, image[0]) +# i += 1 \ No newline at end of file