Overview

This is a small end-to-end computer-vision project: take a photo of a coffee mug and predict its color. The classifier sorts an image into one of four color classes - white, black, blue, and transparent (glass) - and returns a probability for each class. The interesting part is not just the model but the full path from raw images to a live REST prediction API: a convolutional neural network is trained in TensorFlow, exported as a SavedModel, and deployed to Google Cloud Vertex AI as a managed endpoint that answers HTTP requests with class probabilities.

Data & training pipeline

The dataset consists of RGB photographs of mugs organised into four class folders (white = class 0, black = class 1, blue = class 2, transparent = class 3). Images are read with OpenCV, converted to RGB, stacked into NumPy arrays, and rescaled to the [0, 1] range before being fed to the network. The data ships pre-split into a train folder and an eval folder.

There are two training entry points. During development, trainer.task trains on the train split and measures accuracy on the held-out eval split - useful for honest iteration on the architecture. Once the model is finalised, trainer.final_task trains on all available data (train + eval combined) to squeeze out extra performance and exports the fitted model for deployment. Final quality is judged on a hidden holdout set that the model never sees during training.

Data flow diagram: development training uses the train split evaluated on the eval split, while final training uses all data and is scored on a hidden holdout

Two training modes: the development run trains on the train split and evaluates on the eval split, while the final run trains on all available data and is scored on a hidden holdout set.

Model

The core classifier is a convolutional neural network built from scratch. It stacks three convolution + pooling blocks that grow the channel depth while shrinking the spatial resolution, then flattens into a small fully connected head:

  • Conv2D(16, 3x3, ReLU) then MaxPooling - first feature extractor.
  • Conv2D(32, 3x3, ReLU) then MaxPooling - mid-level features.
  • Conv2D(64, 3x3, ReLU) then MaxPooling - higher-level features.
  • Flatten then Dense(128, ReLU) - fully connected head.
  • Dense(num_classes) - one logit per class, turned into a probability distribution over the four mug colors.

The network is compiled with the Adam optimizer and a sparse categorical cross-entropy loss, tracking accuracy. Batch size and epoch count are configurable. A MobileNetV2 transfer-learning variant is also available: it uses an ImageNet-pretrained MobileNetV2 backbone (frozen) followed by GlobalAveragePooling2D and a Dense softmax layer over the four classes - a strong baseline when data is scarce.

The from-scratch stack is defined as a plain Keras Sequential model:

model = Sequential([
    input_layer,
    Conv2D(16, 3, padding='same', activation='relu'),
    MaxPooling2D(),
    Conv2D(32, 3, padding='same', activation='relu'),
    MaxPooling2D(),
    Conv2D(64, 3, padding='same', activation='relu'),
    MaxPooling2D(),
    Flatten(),
    Dense(128, activation='relu'),
    Dense(num_classes)
])

model.compile(
    optimizer='adam',
    loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=False),
    metrics=['accuracy'])

Deployment

Turning the trained network into an API takes a few steps on Google Cloud:

  • Export a SavedModel. Running python -m trainer.final_task trains on the full dataset and writes a TensorFlow SavedModel. The export wraps the network so it accepts a base64-encoded JPEG as a byte string, decodes and preprocesses it inside the graph, and exposes two named outputs: CLASSES (the argmax label) and PROBABILITIES (the softmax vector).
  • Import into the Model Registry. The SavedModel is imported into the Vertex AI Model Registry, using the TensorFlow 2.9 serving container.
  • Deploy to an endpoint. The registered model is deployed to a Vertex AI endpoint in the europe-west1 region (the storage bucket lives in the same region), which spins up managed serving infrastructure behind a stable URL.
  • Query the endpoint. A prediction request sends a JSON payload containing the image bytes; the endpoint returns the predicted class and the per-class probabilities.

A request against the deployed endpoint looks like this:

gcloud ai endpoints predict $ENDPOINT_ID \
    --project=$PROJECT_ID \
    --region=europe-west1 \
    --json-request=test.json

where test.json holds the image as base64 bytes under an instances list.

Example inference

Below is the actual image packed into test.json and sent to the live endpoint: a black mug photographed on a desk. The image is decoded, resized to the network input size, and normalised inside the serving graph before the CNN scores it.

Example input: a black mug on a desk, resized to the CNN input resolution
Input sent to the endpoint

The endpoint echoes back the winning class and the full softmax distribution:

CLASSES  PROBABILITIES
1        [2.06e-12, 1.0, 1.74e-13, 1.29e-32]

The model returns class 1 (black) with essentially full confidence, and near-zero probability on white, blue, and glass.

Result & takeaway

The bar for the task was at least 75% accuracy on a hidden holdout set. The exported artifacts here do not include a saved accuracy log, but the deployed endpoint answers this example with near-certain, correct probabilities, and the from-scratch CNN was enough to clear the target without needing the heavier MobileNetV2 backbone.

The value of the project is the full loop: a small from-scratch network solves the classification task, and Vertex AI handles the unglamorous but essential part of turning a local SavedModel into a scalable, queryable prediction service - send an image, get back four class probabilities. Swapping in the MobileNetV2 backbone is a one-line change if more headroom is needed.