import os from keras.preprocessing.image import ImageDataGenerator from keras.models import Sequential from keras.layers import Conv2D, MaxPooling2D, Flatten, Dense train_data_dir = "garbage_photos" location = os.path.realpath(os.path.join(os.getcwd(), os.path.dirname(__file__))) train_data_dir = os.path.join(location, train_data_dir) input_shape = (150, 150, 3) num_classes = 5 batch_size = 32 epochs = 20 train_datagen = ImageDataGenerator(rescale=1./255) train_generator = train_datagen.flow_from_directory( train_data_dir, target_size=(input_shape[0], input_shape[1]), batch_size=batch_size, class_mode='categorical' ) model = Sequential() model.add(Conv2D(32, (3, 3), activation='relu', input_shape=input_shape)) model.add(MaxPooling2D(pool_size=(2, 2))) model.add(Conv2D(64, (3, 3), activation='relu')) model.add(MaxPooling2D(pool_size=(2, 2))) model.add(Conv2D(128, (3, 3), activation='relu')) model.add(MaxPooling2D(pool_size=(2, 2))) model.add(Flatten()) model.add(Dense(128, activation='relu')) model.add(Dense(num_classes, activation='softmax')) model.compile(optimizer='adam', loss='categorical_crossentropy', metrics=['accuracy']) model.fit(train_generator, epochs=epochs) classes = train_generator.class_indices model.save("neuralModel.h5")