import os import glob import PIL from PIL import Image import tensorflow as tf import pickle from tensorflow import keras from keras import layers from keras.models import Sequential import pathlib class NeuralN: # @staticmethod def predict(self,image): data_dir = pathlib.Path('zdjecia') saved_model_path = pathlib.Path('trained_model.h5') class_names_path = pathlib.Path("class_names.pkl") image_count = sum(len(files) for _, _, files in os.walk(data_dir)) print(image_count) # ORK_ARCHER = list(glob.glob('C:\\mobs_photos\\ORK_ARCHER')) # im = PIL.Image.open(ORK_ARCHER[0]) # im.show() if os.path.exists(saved_model_path): model = tf.keras.models.load_model(saved_model_path) print("Saved model loaded") with open(class_names_path, 'rb') as f: class_names = pickle.load(f) print("Class names loaded.") else: train_ds = tf.keras.utils.image_dataset_from_directory( data_dir, validation_split=0.2, subset="training", seed=123, image_size=(180, 180), batch_size=32) val_ds = tf.keras.utils.image_dataset_from_directory( data_dir, validation_split=0.2, subset="validation", seed=123, image_size=(180, 180), batch_size=32) # test_ds = tf.keras.utils.image_dataset_from_directory( # data_dir, # seed=123, # image_size=(180, 180), # batch_size=32) class_names = train_ds.class_names print(class_names) num_classes = len(class_names) model = Sequential([ layers.Rescaling(1. / 255, input_shape=(180, 180, 3)), layers.Conv2D(16, 3, padding='same', activation='relu'), 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) ]) model.compile(optimizer='adam', loss=tf.keras.losses.SparseCategoricalCrossentropy( from_logits=True), metrics=['accuracy']) model.summary() epochs = 1 history = model.fit( train_ds, validation_data=val_ds, epochs=epochs ) model.save("trained_model.h5") print("Model trained and saved.") with open(class_names_path, 'wb') as f: pickle.dump(train_ds.class_names, f) print("Class names saved.") # loaded_model = tf.keras.models.load_model("trained_model.h5") probability_model = tf.keras.Sequential([model, tf.keras.layers.Softmax()]) #image_path = image image_path = pathlib.Path('zdjecia\ORK_ARCHER\ork_archer (942).jpg') image = Image.open(image_path) # Preprocess the image image = image.resize((180, 180)) # Resize to match the input size of the model image_array = tf.keras.preprocessing.image.img_to_array(image) image_array = image_array / 255.0 # Normalize pixel values # Add an extra dimension to the image array image_array = tf.expand_dims(image_array, 0) # Make the prediction predictions = probability_model.predict(image_array) # Convert the predictions to class labels predicted_label = class_names[predictions[0].argmax()] #actions = { # 'ORK_MELEE': 'fight', # 'ORK_ARCHER': 'change_dir', # 'SAURON': 'change_dir' #} # Get the action for the predicted character #action = actions.get(predicted_label, 'unknown') # Print the predicted label print(predicted_label) return predicted_label#, action