63 lines
1.5 KiB
Python
63 lines
1.5 KiB
Python
|
import tensorflow as tf
|
||
|
|
from tensorflow import keras
|
||
|
|
from keras import layers
|
||
|
|
|
||
|
|
# Load and preprocess the dataset
|
||
|
|
# Assuming you have three folders named 'class1', 'class2', and 'class3'
|
||
|
|
# each containing images of their respective classes
|
||
|
|
|
||
|
|
data_dir = 'Training/'
|
||
|
|
image_size = (100, 100)
|
||
|
|
batch_size = 32
|
||
|
|
|
||
|
|
train_ds = tf.keras.preprocessing.image_dataset_from_directory(
|
||
|
|
data_dir,
|
||
|
|
validation_split=0.2,
|
||
|
|
subset="training",
|
||
|
|
seed=123,
|
||
|
|
image_size=image_size,
|
||
|
|
batch_size=batch_size,
|
||
|
|
)
|
||
|
|
|
||
|
|
val_ds = tf.keras.preprocessing.image_dataset_from_directory(
|
||
|
|
data_dir,
|
||
|
|
validation_split=0.2,
|
||
|
|
subset="validation",
|
||
|
|
seed=123,
|
||
|
|
image_size=image_size,
|
||
|
|
batch_size=batch_size,
|
||
|
|
)
|
||
|
|
|
||
|
|
class_names = train_ds.class_names
|
||
|
|
num_classes = len(class_names)
|
||
|
|
|
||
|
|
|
||
|
|
# Create the model
|
||
|
|
model = keras.Sequential([
|
||
|
|
layers.Rescaling(1./255, input_shape=(100, 100, 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)
|
||
|
|
])
|
||
|
|
|
||
|
|
# Compile the model
|
||
|
|
model.compile(optimizer='adam',
|
||
|
|
loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
|
||
|
|
metrics=['accuracy'])
|
||
|
|
|
||
|
|
# Train the model
|
||
|
|
epochs = 10
|
||
|
|
model.fit(
|
||
|
|
train_ds,
|
||
|
|
validation_data=val_ds,
|
||
|
|
epochs=epochs
|
||
|
|
)
|
||
|
|
|
||
|
|
# Save the trained model
|
||
|
|
model.save('trained_model')
|