Saltar navegación

Session 2: Preprocessing and Training the Model

Introdución

📌 Objective: Train a simple CNN model to classify objects.

  • Brief explanation of how CNNs work.
  • Using Jupyter Notebook / Google Colab to prepare data:
    • Resizing and normalizing images.
    • Creating training and validation sets.
  • Building and training a CNN model using TensorFlow/Keras.
  • Evaluating the model and fine-tuning hyperparameters.

💡 By the end of this session, participants will have a trained model ready for deployment.

Hands-on

📂 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")

Feito con eXeLearning (Nova xanela)