27/10/2025, 13:03 Image classification | TensorFlow Core
Image classification
Run in Google Colab ([Link]
View source on GitHub ([Link]
Download notebook ([Link]
This tutorial shows how to classify images of flowers using a [Link]
([Link] model and load data using
[Link].image_dataset_from_directory
([Link]
It demonstrates the following concepts:
Efficiently loading a dataset off disk.
Identifying overfitting and applying techniques to mitigate it, including data
augmentation and dropout.
This tutorial follows a basic machine learning workflow:
1. Examine and understand data
2. Build an input pipeline
3. Build the model
4. Train the model
5. Test the model
6. Improve the model and repeat the process
In addition, the notebook demonstrates how to convert a saved model
([Link] to a TensorFlow Lite
([Link] model for on-device machine learning on mobile, embedded,
and IoT devices.
Setup
Import TensorFlow and other necessary libraries:
[Link] 1/16
27/10/2025, 13:03 Image classification | TensorFlow Core
import [Link] as plt
import numpy as np
import PIL
import tensorflow as tf
from tensorflow import keras
from [Link] import layers
from [Link] import Sequential
Download and explore the dataset
This tutorial uses a dataset of about 3,700 photos of flowers. The dataset contains five sub-
directories, one per class:
flower_photo/
daisy/
dandelion/
roses/
sunflowers/
tulips/
import pathlib
dataset_url = "[Link]
data_dir = [Link].get_file('flower_photos.tar', origin=dataset_url, ex
data_dir = [Link](data_dir).with_suffix('')
After downloading, you should now have a copy of the dataset available. There are 3,670
total images:
image_count = len(list(data_dir.glob('*/*.jpg')))
print(image_count)
Here are some roses:
[Link] 2/16
27/10/2025, 13:03 Image classification | TensorFlow Core
roses = list(data_dir.glob('roses/*'))
[Link](str(roses[0]))
[Link](str(roses[1]))
And some tulips:
tulips = list(data_dir.glob('tulips/*'))
[Link](str(tulips[0]))
[Link](str(tulips[1]))
Load data using a Keras utility
Next, load these images off disk using the helpful
[Link].image_dataset_from_directory
([Link]
utility. This will take you from a directory of images on disk to a [Link]
([Link] in just a couple lines of code. If you
like, you can also write your own data loading code from scratch by visiting the Load and
preprocess images ([Link] tutorial.
Create a dataset
Define some parameters for the loader:
batch_size = 32
img_height = 180
img_width = 180
[Link] 3/16
27/10/2025, 13:03 Image classification | TensorFlow Core
It's good practice to use a validation split when developing your model. Use 80% of the
images for training and 20% for validation.
train_ds = [Link].image_dataset_from_directory(
data_dir,
validation_split=0.2,
subset="training",
seed=123,
image_size=(img_height, img_width),
batch_size=batch_size)
val_ds = [Link].image_dataset_from_directory(
data_dir,
validation_split=0.2,
subset="validation",
seed=123,
image_size=(img_height, img_width),
batch_size=batch_size)
You can find the class names in the class_names attribute on these datasets. These
correspond to the directory names in alphabetical order.
class_names = train_ds.class_names
print(class_names)
Visualize the data
Here are the first nine images from the training dataset:
import [Link] as plt
[Link](figsize=(10, 10))
for images, labels in train_ds.take(1):
for i in range(9):
ax = [Link](3, 3, i + 1)
[Link](images[i].numpy().astype("uint8"))
[Link] 4/16
27/10/2025, 13:03 Image classification | TensorFlow Core
[Link](class_names[labels[i]])
[Link]("off")
You will pass these datasets to the Keras [Link]
([Link] method for training later in this
tutorial. If you like, you can also manually iterate over the dataset and retrieve batches of
images:
for image_batch, labels_batch in train_ds:
print(image_batch.shape)
print(labels_batch.shape)
break
The image_batch is a tensor of the shape (32, 180, 180, 3). This is a batch of 32
images of shape 180x180x3 (the last dimension refers to color channels RGB). The
label_batch is a tensor of the shape (32,), these are corresponding labels to the 32
images.
You can call .numpy() on the image_batch and labels_batch tensors to convert them to a
[Link].
Configure the dataset for performance
Make sure to use buffered prefetching, so you can yield data from disk without having I/O
become blocking. These are two important methods you should use when loading data:
[Link] ([Link] keeps
the images in memory after they're loaded off disk during the first epoch. This will
ensure the dataset does not become a bottleneck while training your model. If your
dataset is too large to fit into memory, you can also use this method to create a
performant on-disk cache.
[Link] ([Link]
overlaps data preprocessing and model execution while training.
Interested readers can learn more about both methods, as well as how to cache data to
disk in the Prefetching section of the Better performance with the [Link] API
([Link] guide.
[Link] 5/16
27/10/2025, 13:03 Image classification | TensorFlow Core
AUTOTUNE = [Link]
train_ds = train_ds.cache().shuffle(1000).prefetch(buffer_size=AUTOTUNE)
val_ds = val_ds.cache().prefetch(buffer_size=AUTOTUNE)
Standardize the data
The RGB channel values are in the [0, 255] range. This is not ideal for a neural network; in
general you should seek to make your input values small.
Here, you will standardize values to be in the [0, 1] range by using
[Link]
([Link]
normalization_layer = [Link](1./255)
There are two ways to use this layer. You can apply it to the dataset by calling [Link]
([Link]
normalized_ds = train_ds.map(lambda x, y: (normalization_layer(x), y))
image_batch, labels_batch = next(iter(normalized_ds))
first_image = image_batch[0]
# Notice the pixel values are now in `[0,1]`.
print([Link](first_image), [Link](first_image))
Or, you can include the layer inside your model definition, which can simplify deployment.
Use the second approach here.
Note: You previously resized images using the image_size argument of
[Link].image_dataset_from_directory
([Link]
If you want to include the resizing logic in your model as well, you can use the
[Link] ([Link]
layer.
[Link] 6/16
27/10/2025, 13:03 Image classification | TensorFlow Core
A basic Keras model
Create the model
The Keras Sequential ([Link] model consists
of three convolution blocks ([Link].Conv2D
([Link] with a max pooling layer
([Link].MaxPooling2D
([Link] in each of them. There's
a fully-connected layer ([Link]
([Link] with 128 units on top of it
that is activated by a ReLU activation function ('relu'). This model has not been tuned for
high accuracy; the goal of this tutorial is to show a standard approach.
num_classes = len(class_names)
model = Sequential([
[Link](1./255, input_shape=(img_height, img_width, 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(),
[Link](),
[Link](128, activation='relu'),
[Link](num_classes)
])
Compile the model
For this tutorial, choose the [Link]
([Link] optimizer and
[Link]
([Link] loss
function. To view training and validation accuracy for each training epoch, pass the metrics
argument to [Link]
([Link]
[Link] 7/16
27/10/2025, 13:03 Image classification | TensorFlow Core
[Link](optimizer='adam',
loss=[Link](from_logits=T
metrics=['accuracy'])
Model summary
View all the layers of the network using the Keras [Link]
([Link] method:
[Link]()
Train the model
Train the model for 10 epochs with the Keras [Link]
([Link] method:
epochs=10
history = [Link](
train_ds,
validation_data=val_ds,
epochs=epochs
)
Visualize training results
Create plots of the loss and accuracy on the training and validation sets:
acc = [Link]['accuracy']
val_acc = [Link]['val_accuracy']
loss = [Link]['loss']
val_loss = [Link]['val_loss']
epochs_range = range(epochs)
[Link] 8/16
27/10/2025, 13:03 Image classification | TensorFlow Core
[Link](figsize=(8, 8))
[Link](1, 2, 1)
[Link](epochs_range, acc, label='Training Accuracy')
[Link](epochs_range, val_acc, label='Validation Accuracy')
[Link](loc='lower right')
[Link]('Training and Validation Accuracy')
[Link](1, 2, 2)
[Link](epochs_range, loss, label='Training Loss')
[Link](epochs_range, val_loss, label='Validation Loss')
[Link](loc='upper right')
[Link]('Training and Validation Loss')
[Link]()
The plots show that training accuracy and validation accuracy are off by large margins, and
the model has achieved only around 60% accuracy on the validation set.
The following tutorial sections show how to inspect what went wrong and try to increase
the overall performance of the model.
Overfitting
In the plots above, the training accuracy is increasing linearly over time, whereas validation
accuracy stalls around 60% in the training process. Also, the difference in accuracy between
training and validation accuracy is noticeable—a sign of overfitting
([Link]
When there are a small number of training examples, the model sometimes learns from
noises or unwanted details from training examples—to an extent that it negatively impacts
the performance of the model on new examples. This phenomenon is known as overfitting.
It means that the model will have a difficult time generalizing on a new dataset.
There are multiple ways to fight overfitting in the training process. In this tutorial, you'll use
data augmentation and add dropout to your model.
Data augmentation
Overfitting generally occurs when there are a small number of training examples. Data
augmentation ([Link] takes the
[Link] 9/16
27/10/2025, 13:03 Image classification | TensorFlow Core
approach of generating additional training data from your existing examples by augmenting
them using random transformations that yield believable-looking images. This helps expose
the model to more aspects of the data and generalize better.
You will implement data augmentation using the following Keras preprocessing layers:
[Link]
([Link]
[Link]
([Link] and
[Link]
([Link] These can be
included inside your model like other layers, and run on the GPU.
data_augmentation = [Link](
[
[Link]("horizontal",
input_shape=(img_height,
img_width,
3)),
[Link](0.1),
[Link](0.1),
]
)
Visualize a few augmented examples by applying data augmentation to the same image
several times:
[Link](figsize=(10, 10))
for images, _ in train_ds.take(1):
for i in range(9):
augmented_images = data_augmentation(images)
ax = [Link](3, 3, i + 1)
[Link](augmented_images[0].numpy().astype("uint8"))
[Link]("off")
You will add data augmentation to your model before training in the next step.
Dropout
[Link] 10/16
27/10/2025, 13:03 Image classification | TensorFlow Core
Another technique to reduce overfitting is to introduce dropout
([Link] regularization to
the network.
When you apply dropout to a layer, it randomly drops out (by setting the activation to zero) a
number of output units from the layer during the training process. Dropout takes a
fractional number as its input value, in the form such as 0.1, 0.2, 0.4, etc. This means
dropping out 10%, 20% or 40% of the output units randomly from the applied layer.
Create a new neural network with [Link]
([Link] before training it using the
augmented images:
model = Sequential([
data_augmentation,
[Link](1./255),
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(),
[Link](0.2),
[Link](),
[Link](128, activation='relu'),
[Link](num_classes, name="outputs")
])
Compile and train the model
[Link](optimizer='adam',
loss=[Link](from_logits=T
metrics=['accuracy'])
[Link]()
[Link] 11/16
27/10/2025, 13:03 Image classification | TensorFlow Core
epochs = 15
history = [Link](
train_ds,
validation_data=val_ds,
epochs=epochs
)
Visualize training results
After applying data augmentation and [Link]
([Link] there is less overfitting than
before, and training and validation accuracy are closer aligned:
acc = [Link]['accuracy']
val_acc = [Link]['val_accuracy']
loss = [Link]['loss']
val_loss = [Link]['val_loss']
epochs_range = range(epochs)
[Link](figsize=(8, 8))
[Link](1, 2, 1)
[Link](epochs_range, acc, label='Training Accuracy')
[Link](epochs_range, val_acc, label='Validation Accuracy')
[Link](loc='lower right')
[Link]('Training and Validation Accuracy')
[Link](1, 2, 2)
[Link](epochs_range, loss, label='Training Loss')
[Link](epochs_range, val_loss, label='Validation Loss')
[Link](loc='upper right')
[Link]('Training and Validation Loss')
[Link]()
Predict on new data
Use your model to classify an image that wasn't included in the training or validation sets.
[Link] 12/16
27/10/2025, 13:03 Image classification | TensorFlow Core
Note: Data augmentation and dropout layers are inactive at inference time.
sunflower_url = "[Link]
sunflower_path = [Link].get_file('Red_sunflower', origin=sunflower_url
img = [Link].load_img(
sunflower_path, target_size=(img_height, img_width)
)
img_array = [Link].img_to_array(img)
img_array = tf.expand_dims(img_array, 0) # Create a batch
predictions = [Link](img_array)
score = [Link](predictions[0])
print(
"This image most likely belongs to {} with a {:.2f} percent confidence."
.format(class_names[[Link](score)], 100 * [Link](score))
)
Use TensorFlow Lite
TensorFlow Lite is a set of tools that enables on-device machine learning by helping
developers run their models on mobile, embedded, and edge devices.
Convert the Keras Sequential model to a TensorFlow Lite model
To use the trained model with on-device applications, first convert it
([Link] to a smaller and more efficient model format
called a TensorFlow Lite ([Link] model.
In this example, take the trained Keras Sequential model and use
[Link].from_keras_model
([Link] to generate
a TensorFlow Lite ([Link] model:
# Convert the model.
converter = [Link].from_keras_model(model)
tflite_model = [Link]()
[Link] 13/16
27/10/2025, 13:03 Image classification | TensorFlow Core
# Save the model.
with open('[Link]', 'wb') as f:
[Link](tflite_model)
The TensorFlow Lite model you saved in the previous step can contain several function
signatures. The Keras model converter API uses the default signature automatically. Learn
more about TensorFlow Lite signatures ([Link]
Run the TensorFlow Lite model
You can access the TensorFlow Lite saved model signatures in Python via the
[Link] ([Link] class.
Load the model with the Interpreter:
TF_MODEL_FILE_PATH = '[Link]' # The default path to the saved TensorFlow
interpreter = [Link](model_path=TF_MODEL_FILE_PATH)
Print the signatures from the converted model to obtain the names of the inputs (and
outputs):
interpreter.get_signature_list()
In this example, you have one default signature called serving_default. In addition, the
name of the 'inputs' is 'sequential_1_input', while the 'outputs' are called
'outputs'. You can look up these first and last Keras layer names when running
[Link], as demonstrated earlier in this tutorial.
Now you can test the loaded TensorFlow Model by performing inference on a sample image
with [Link].get_signature_runner
([Link] by passing the
signature name as follows:
classify_lite = interpreter.get_signature_runner('serving_default')
classify_lite
[Link] 14/16
27/10/2025, 13:03 Image classification | TensorFlow Core
Similar to what you did earlier in the tutorial, you can use the TensorFlow Lite model to
classify images that weren't included in the training or validation sets.
You have already tensorized that image and saved it as img_array. Now, pass it to the first
argument (the name of the 'inputs') of the loaded TensorFlow Lite model
(predictions_lite), compute softmax activations, and then print the prediction for the
class with the highest computed probability.
predictions_lite = classify_lite(sequential_1_input=img_array)['outputs']
score_lite = [Link](predictions_lite)
print(
"This image most likely belongs to {} with a {:.2f} percent confidence."
.format(class_names[[Link](score_lite)], 100 * [Link](score_lite))
)
The prediction generated by the lite model should be almost identical to the predictions
generated by the original model:
print([Link]([Link](predictions - predictions_lite)))
Of the five classes—'daisy', 'dandelion', 'roses', 'sunflowers', and 'tulips'—the
model should predict the image belongs to sunflowers, which is the same result as before
the TensorFlow Lite conversion.
Next steps
This tutorial showed how to train a model for image classification, test it, convert it to the
TensorFlow Lite format for on-device applications (such as an image classification app),
and perform inference with the TensorFlow Lite model with the Python API.
You can learn more about TensorFlow Lite through tutorials
([Link] and guides ([Link]
[Link] 15/16
27/10/2025, 13:03 Image classification | TensorFlow Core
Except as otherwise noted, the content of this page is licensed under the Creative Commons Attribution 4.0
License ([Link] and code samples are licensed under the Apache
2.0 License ([Link] For details, see the Google Developers Site
Policies ([Link] Java is a registered trademark of Oracle and/or its
affiliates.
Last updated 2024-04-03 UTC.
[Link] 16/16