📂 Step 1: Load and Preprocess the Dataset
Now that we have collected the images, we need to prepare them for training. We will use TensorFlow's ImageDataGenerator
to automatically resize, normalize, and split the images into training and validation sets. This ensures the model learns from a diverse set of images while preventing overfitting.
import tensorflow as tf
from tensorflow.keras.preprocessing.image import ImageDataGenerator
# Path to the dataset folder
dataset_path = "dataset/"
# Create a data generator to preprocess images
datagen = ImageDataGenerator(
rescale=1.0/255.0, # Normalize pixel values (0-1)
validation_split=0.2 # Use 20% of images for validation
)
# Load training and validation sets
train_generator = datagen.flow_from_directory(
dataset_path,
target_size=(64, 64), # Resize images to 64x64
batch_size=32,
class_mode='categorical',
subset='training'
)
val_generator = datagen.flow_from_directory(
dataset_path,
target_size=(64, 64),
batch_size=32,
class_mode='categorical',
subset='validation'
)
# Get class names
class_names = list(train_generator.class_indices.keys())
print("Classes:", class_names)
📂 Step 2: Build and Train the CNN Model
Now we will create a Convolutional Neural Network (CNN) using TensorFlow and Keras. The CNN consists of multiple layers that help extract meaningful features from the images, such as edges, shapes, and textures. The model will be trained on our dataset and will learn to classify waste materials into different categories.
from tensorflow.keras import models, layers
# Create a simple CNN model
model = models.Sequential([
layers.Conv2D(32, (3,3), activation='relu', input_shape=(64, 64, 3)),
layers.MaxPooling2D(2,2),
layers.Conv2D(64, (3,3), activation='relu'),
layers.MaxPooling2D(2,2),
layers.Flatten(),
layers.Dense(128, activation='relu'),
layers.Dense(len(class_names), activation='softmax') # Output layer matches number of classes
])
# Compile the model
model.compile(optimizer='adam',
loss='categorical_crossentropy',
metrics=['accuracy'])
# Train the model
history = model.fit(train_generator, validation_data=val_generator, epochs=10)
# Save the trained model
model.save("recycling_model.h5")
print("Model saved as recycling_model.h5")