import tensorflow as tf
from tensorflow.keras import layers, models
from tensorflow.keras.preprocessing.image import ImageDataGenerator
# Data preparation
train_dir = 'path/to/training/data'
validation_dir = 'path/to/validation/data'
# Create data generators with data augmentation for training
train_datagen = ImageDataGenerator(
rescale=1.0/255,
rotation_range=40,
width_shift_range=0.2,
height_shift_range=0.2,
shear_range=0.2,
zoom_range=0.2,
horizontal_flip=True,
fill_mode='nearest'
)
validation_datagen = ImageDataGenerator(rescale=1.0/255)
train_generator = train_datagen.flow_from_directory(
train_dir,
target_size=(150, 150),
batch_size=20,
class_mode='categorical'
)
validation_generator = validation_datagen.flow_from_directory(
validation_dir,
target_size=(150, 150),
batch_size=20,
class_mode='categorical'
)
# Build the model
model = models.Sequential([
layers.Conv2D(32, (3, 3), activation='relu', input_shape=(150, 150, 3)),
layers.MaxPooling2D((2, 2)),
layers.Conv2D(64, (3, 3), activation='relu'),
layers.MaxPooling2D((2, 2)),
layers.Conv2D(128, (3, 3), activation='relu'),
layers.MaxPooling2D((2, 2)),
layers.Flatten(),
layers.Dense(512, activation='relu'),
layers.Dense(3, activation='softmax')
])
# Compile the model
model.compile(optimizer='adam',
loss='categorical_crossentropy',
metrics=['accuracy'])
# Train the model
history = model.fit(
train_generator,
steps_per_epoch=100,
epochs=10,
validation_data=validation_generator,
validation_steps=50
)
# Save the model
model.save('image_classifier_model.h5')
# Load the model
loaded_model = tf.keras.models.load_model('image_classifier_model.h5')
# Make predictions on new images
new_image = tf.keras.preprocessing.image.load_img('path/to/new/image.jpg', target_size=(150, 150))
new_image = tf.keras.preprocessing.image.img_to_array(new_image)
new_image = new_image.reshape((1, 150, 150, 3))
new_image = new_image / 255.0 # Normalize pixel values
predictions = loaded_model.predict(new_image)