Using Your Own Models

PeekingDuck offers pre-trained model nodes that can be used to tackle a wide variety of problems, but you may need to train your own model on a custom dataset sometimes. This tutorial will show you how to package your model into a custom model node, and use it with PeekingDuck. We will be tackling a manufacturing use case here - classifying images of steel castings into “defective” or “normal” classes.

Casting is a manufacturing process in which a material such as metal in liquid form is poured into a mold and allowed to solidify. The solidified result is also called a casting. Sometimes, defective castings are produced, and quality assurance departments are responsible for preventing defective pieces from being used downstream. As inspections are usually done manually, this adds a significant amount of time and cost, and thus there is an incentive to automate this process.

The images of castings used in this tutorial are the front faces of steel pump impellers. From the comparison below, it can be seen that the defective casting has a rough, uneven edges while the normal casting has smooth edges.

Normal casting compared to defective casting

Normal Casting Compared to Defective Casting

Model Training

PeekingDuck is designed for model inference rather than model training. This optional section shows how a simple Convolutional Neural Network (CNN) model can be trained separately from the PeekingDuck framework. If you have already trained your own model, the following section describes how you can convert it to a custom model node, and use it within PeekingDuck for inference.

Setting Up

Install the following prerequisite for visualization.

> conda install matplotlib

Create the following project folder:

Terminal Session

[~user] > mkdir castings_project
[~user] > cd castings_project

Download the castings dataset and unzip the file to the castings_project folder.


The castings dataset used in this example is modified from the original dataset from Kaggle.

You should have the following directory structure at this point:

castings_project/ ⠀
└── castings_data/ ⠀
      ├── inspection/ ⠀
      ├── train/ ⠀
      └── validation/

Update Training Script

Create an empty file within the castings_project folder, and update it with the following code:

Show/Hide Code for

  2Script to train a classification model on images, save the model, and plot the training results
  4Adapted from:
  7import pathlib
  8from typing import List, Tuple
 10import matplotlib.pyplot as plt
 11import tensorflow as tf
 12from tensorflow.keras import layers
 13from tensorflow.keras.models import Sequential
 14from tensorflow.keras.layers.experimental.preprocessing import Rescaling
 16# setup global constants
 17DATA_DIR = "./castings_data"
 18WEIGHTS_DIR = "./weights"
 19RESULTS = "training_results.png"
 20EPOCHS = 10
 21BATCH_SIZE = 32
 22IMG_HEIGHT = 180
 23IMG_WIDTH = 180
 26def prepare_data() -> Tuple[,, List[str]]:
 27   """
 28   Generate training and validation datasets from a folder of images.
 30   Returns:
 31      train_ds ( Training dataset.
 32      val_ds ( Validation dataset.
 33      class_names (List[str]): Names of all classes to be classified.
 34   """
 36   train_dir = pathlib.Path(DATA_DIR, "train")
 37   validation_dir = pathlib.Path(DATA_DIR, "validation")
 39   train_ds = tf.keras.preprocessing.image_dataset_from_directory(
 40      train_dir,
 41      image_size=(IMG_HEIGHT, IMG_WIDTH),
 42      batch_size=BATCH_SIZE,
 43   )
 45   val_ds = tf.keras.preprocessing.image_dataset_from_directory(
 46      validation_dir,
 47      image_size=(IMG_HEIGHT, IMG_WIDTH),
 48      batch_size=BATCH_SIZE,
 49   )
 51   class_names = train_ds.class_names
 53   return train_ds, val_ds, class_names
 56def train_and_save_model(
 57   train_ds:, val_ds:, class_names: List[str]
 58) -> tf.keras.callbacks.History:
 59   """
 60   Train and save a classification model on the provided data.
 62   Args:
 63      train_ds ( Training dataset.
 64      val_ds ( Validation dataset.
 65      class_names (List[str]): Names of all classes to be classified.
 67   Returns:
 68      history (tf.keras.callbacks.History): A History object containing recorded events from
 69               model training.
 70   """
 72   num_classes = len(class_names)
 74   model = Sequential(
 75      [
 76            Rescaling(1.0 / 255, input_shape=(IMG_HEIGHT, IMG_WIDTH, 3)),
 77            layers.Conv2D(16, 3, padding="same", activation="relu"),
 78            layers.MaxPooling2D(),
 79            layers.Conv2D(32, 3, padding="same", activation="relu"),
 80            layers.MaxPooling2D(),
 81            layers.Conv2D(64, 3, padding="same", activation="relu"),
 82            layers.MaxPooling2D(),
 83            layers.Dropout(0.2),
 84            layers.Flatten(),
 85            layers.Dense(128, activation="relu"),
 86            layers.Dense(num_classes),
 87      ]
 88   )
 90   model.compile(
 91      optimizer="adam",
 92      loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
 93      metrics=["accuracy"],
 94   )
 96   print(model.summary())
 97   history =, validation_data=val_ds, epochs=EPOCHS)
100   return history
103def plot_training_results(history: tf.keras.callbacks.History) -> None:
104   """
105   Plot training and validation accuracy and loss curves, and save the plot.
107   Args:
108      history (tf.keras.callbacks.History): A History object containing recorded events from
109               model training.
110   """
111   acc = history.history["accuracy"]
112   val_acc = history.history["val_accuracy"]
113   loss = history.history["loss"]
114   val_loss = history.history["val_loss"]
115   epochs_range = range(EPOCHS)
117   plt.figure(figsize=(16, 8))
118   plt.subplot(1, 2, 1)
119   plt.plot(epochs_range, acc, label="Training Accuracy")
120   plt.plot(epochs_range, val_acc, label="Validation Accuracy")
121   plt.legend(loc="lower right")
122   plt.title("Training and Validation Accuracy")
124   plt.subplot(1, 2, 2)
125   plt.plot(epochs_range, loss, label="Training Loss")
126   plt.plot(epochs_range, val_loss, label="Validation Loss")
127   plt.legend(loc="upper right")
128   plt.title("Training and Validation Loss")
129   plt.savefig(RESULTS)
132if __name__ == "__main__":
133   train_ds, val_ds, class_names = prepare_data()
134   history = train_and_save_model(train_ds, val_ds, class_names)
135   plot_training_results(history)

Training the Model

Train the model by running the following command.

Terminal Session

[~user/castings_project] > python


For macOS Apple Silicon, the above code only works on macOS 12.x Monterey with the latest tensorflow-macos and tensorflow-metal versions. It will crash on macOS 11.x Big Sur due to bugs in the outdated tensorflow versions.

The model will be trained for 10 epochs, and when training is completed, a new weights folder and training_results.png will be created:

castings_project/ ⠀
├── training_results.png
├── castings_data/ ⠀
│     ├── inspection/ ⠀
│     ├── train/ ⠀
│     └── validation/ ⠀
└── weights/ ⠀
      ├── keras_metadata.pb
      ├── saved_model.pb
      ├── assets/ ⠀
      └── variables/

The plots from training_results.png shown below indicate that the model has performed well on the validation dataset, and we are ready to create a custom model node from it.

Model training results

Model Training Results

Using Your Trained Model with PeekingDuck

This section will show you how to convert your trained model into a custom PeekingDuck node, and give an example of how you can integrate this node in a PeekingDuck pipeline. It assumes that you are already familiar with the process of creating custom nodes, covered in the earlier custom node tutorial.

Converting to a Custom Model Node

First, let’s create a new PeekingDuck project within the existing castings_project folder.

Terminal Session

[~user/castings_project] > peekingduck init

Next, we’ll use the peekingduck create-node command to create a custom node:

Terminal Session

[~user/castings_project] > peekingduck create-node
Creating new custom node…
Enter node directory relative to ~user/castings_project [src/custom_nodes]:
Select node type (input, augment, model, draw, dabble, output): model
Enter node name [my_custom_node]: casting_classifier

Node directory: ~user/castings_project/src/custom_nodes
Node type: model
Node name: casting_classifier

Creating the following files:
Config file: ~user/castings_project/src/custom_nodes/configs/model/casting_classifier.yml
Script file: ~user/castings_project/src/custom_nodes/model/
Proceed? [Y/n]:
Created node!

The castings_project folder structure should now look like this:

castings_project/ ⠀
├── pipeline_config.yml
├── training_results.png
├── castings_data/ ⠀
│     ├── inspection/ ⠀
│     ├── train/ ⠀
│     └── validation/ ⠀
├── src/ ⠀
│     └── custom_nodes/ ⠀
│           ├── configs/ ⠀
│           │     └── model/ ⠀
│           │           └── casting_classifier.yml
│           └── model/ ⠀
│                 └──
└── weights/ ⠀
      ├── keras_metadata.pb
      ├── saved_model.pb
      ├── assets/ ⠀
      └── variables/

castings_project now contains two files that we need to modify to implement our custom node.

  1. src/custom_nodes/configs/model/casting_classifier.yml:

    casting_classifier.yml updated content:

    1input: ["img"]
    2output: ["pred_label", "pred_score"]
    4weights_parent_dir: weights
    5class_label_map: {0: "defective", 1: "normal"}
  2. src/custom_nodes/model/ updated content:

    Show/Hide Code for

     2Casting classification model.
     5from typing import Any, Dict
     7import cv2
     8import numpy as np
     9import tensorflow as tf
    11from peekingduck.pipeline.nodes.node import AbstractNode
    13IMG_HEIGHT = 180
    14IMG_WIDTH = 180
    17class Node(AbstractNode):
    18   """Initializes and uses a CNN to predict if an image frame shows a normal
    19   or defective casting.
    20   """
    22   def __init__(self, config: Dict[str, Any] = None, **kwargs: Any) -> None:
    23      super().__init__(config, node_path=__name__, **kwargs)
    24      self.model = tf.keras.models.load_model(self.weights_parent_dir)
    26   def run(self, inputs: Dict[str, Any]) -> Dict[str, Any]:
    27      """Reads the image input and returns the predicted class label and
    28      confidence score.
    30      Args:
    31            inputs (dict): Dictionary with key "img".
    33      Returns:
    34            outputs (dict): Dictionary with keys "pred_label" and "pred_score".
    35      """
    36      img = cv2.cvtColor(inputs["img"], cv2.COLOR_BGR2RGB)
    37      img = cv2.resize(img, (IMG_WIDTH, IMG_HEIGHT))
    38      img = np.expand_dims(img, axis=0)
    39      predictions = self.model.predict(img)
    40      score = tf.nn.softmax(predictions[0])
    42      return {
    43            "pred_label": self.class_label_map[np.argmax(score)],
    44            "pred_score": 100.0 * np.max(score),
    45      }

The custom node takes in the built-in PeekingDuck img data type, makes a prediction based on the image, and produces two custom data types: pred_label, the predicted label (“defective” or “normal”); and pred_score, which is the confidence score of the prediction.

Using the Classifier in a PeekingDuck Pipeline

We’ll now pair this custom node with other PeekingDuck nodes to build a complete solution. Imagine an automated inspection system like the one shown below, where the castings are placed on a conveyor belt and a camera takes a picture of each casting and sends it to the PeekingDuck pipeline for prediction. A report showing the predicted result for each casting is produced, and the quality inspector can use it for further analysis.

Vision Based Inspection of Conveyed Objects

Vision Based Inspection of Conveyed Objects (Source: ScienceDirect)

Edit the pipeline_config.yml file to use the input.visual node to read in the images, and the output.csv_writer node to produce the report. We will test our solution on the 10 casting images in castings_data/inspection, where each image’s filename is a unique casting ID such as 28_4137.jpeg.


pipeline_config.yml updated content:

2- input.visual:
3   source: castings_data/inspection
4- custom_nodes.model.casting_classifier
5- output.csv_writer:
6   stats_to_track: ["filename", "pred_label", "pred_score"]
7   file_path: casting_predictions.csv
8   logging_interval: 0
Line 2 input.visual: tells PeekingDuck to load the images from castings_data/inspection.
Line 4 Calls the custom model node that you have just created.
Line 5 output.csv_writer: produces the report for the quality inspector in a CSV file castings_predictions_DDMMYY-hh-mm-ss.csv (time stamp appended to file_path). This node receives the filename data type from input.visual, the custom data types pred_label and pred_score from the custom model node, and writes them to the columns of the CSV file.

Run the above with the command peekingduck run.

Open the created CSV file and you would see the following results. Half of the castings have been predicted as defective with high confidence scores. As the file name of each image is its unique casting ID, the quality inspector would be able to check the results with the actual castings if needed.

Casting prediction results

Casting Prediction Results

To visualize the predictions alongside the casting images, create an empty Python script named, and update it with the following code:

Show/Hide Code for

 2Script to visualize the prediction results alongside the casting images
 5import csv
 7import cv2
 8import matplotlib.pyplot as plt
10CSV_FILE = "casting_predictions_280422-11-50-30.csv"  # change file name accordingly
11INSPECTION_IMGS_DIR = "castings_data/inspection/"
12RESULTS_FILE = "inspection_results.png"
14fig, axs = plt.subplots(2, 5, figsize=(50, 20))
16with open(CSV_FILE) as csv_file:
17   csv_reader = csv.reader(csv_file, delimiter=",")
18   next(csv_reader, None)
19   for i, row in enumerate(csv_reader):
20      # csv columns follow this order: 'Time', 'filename', 'pred_label', 'pred_score'
21      image_path = INSPECTION_IMGS_DIR + row[1]
22      image_orig = cv2.imread(image_path)
23      image_orig = cv2.cvtColor(image_orig, cv2.COLOR_BGR2RGB)
25      row_idx = 0 if i < 5 else 1
26      axs[row_idx][i % 5].imshow(image_orig)
27      axs[row_idx][i % 5].set_title(row[1] + " - " + row[2], fontsize=35)
28      axs[row_idx][i % 5].axis("off")

In Line 10, replace the name of CSV_FILE with the name of the CSV file produced on your system, as a timestamp would have been appended to the file name.

Run the following command to visualize the results.

Terminal Session

[~user/castings_project] > python

An inspection_results.png would be created, as shown below. The top row of castings are clearly defective, as they have rough, uneven edges, while the bottom row of castings look normal. Therefore, the prediction results are accurate for this batch of inspected castings. The quality inspector can provide feedback to the manufacturing team to further investigate the defective castings based on the casting IDs.

Casting prediction visualization

Casting Prediction Visualization

This concludes the guided example on using your own custom models.

Custom Object Detection Models

The previous example was centered on the task of image classification. Object detection is another common task in Computer Vision. PeekingDuck offers several pre-trained object detection model nodes which can detect up to 80 different types of objects, such as persons, cars, and dogs, just to name a few. For the complete list of detectable objects, refer to the Object Detection IDs page. Quite often, you may need to train a custom object detection model on your own dataset, such as defects on a printed circuit board (PCB) as shown below. This section discusses some important considerations for the object detection task, supplementing the guided example above.

Object detection of defects on PCB

Object Detection of Defects on PCB (Source: The Institution of Engineering and Technology)

PeekingDuck’s object detection model nodes conventionally receive the img data type, and produce the bboxes, bbox_labels, and bbox_scores data types. An example of this can be seen in the API documentation for a node such as model.efficientdet. We strongly recommend keeping to these data type conventions for your custom object detection node, ensuring that they adhere to the described format, e.g. img is in BGR format, and bboxes is a NumPy array of a certain shape.

This allows you to leverage on PeekingDuck’s ecosystem of existing nodes. For example, by ensuring that your custom model node receives img in the correct format, you are able to use PeekingDuck’s input.visual node, which can read from multiple visual sources such as a folder of images or videos, an online cloud source, or a CCTV/webcam live feed. By ensuring that your custom model node produces bboxes and bbox_labels in the correct format, you are able to use PeekingDuck’s draw.bbox node to draw bounding boxes and associated labels around the detected objects.

By doing so, you would have saved a significant amount of development time, and can focus more on developing and finetuning your custom object detection model. This was just a simple example, and you can find out more about PeekingDuck’s nodes from our API Documentation, and PeekingDuck’s built-in data types from our Glossary.