How to Train a Custom Model in ml5.js?

Last Updated : 27 Aug, 2024

In machine learning, custom model training helps programmers create models that are optimized for tasks, enhancing the precision of AI systems. This article guides you to training a custom model using ml5.js demonstrating how to capture images using a webcam.

Approach

Set up the basic structure of the webpage.
Include the ml5.js library and p5.js library for handling graphics.
Create a <video> element to capture webcam video and a <canvas> element to display the webcam feed and interact with the model.
Apply These Methods:
- setup(): Initializes the canvas, video capture, and the feature classifier model. Sets up the buttons for training and saving the model.
- modelReady(): Logs a message when the model is loaded and ready to use.
- videoReady(): Logs a message when the webcam video is loaded.
- whileTraining(loss): Displays the training progress, showing the loss as the model trains.
- trainModel(): Starts training the model with the captured images.
- saveModel(): Saves the trained model to the local storage.
- draw(): Continuously draws the video feed on the canvas.

Example: This code sets up a webpage that uses p5.js and ml5.js to display a webcam feed and provides buttons for training and saving a custom model.

HTML

<!DOCTYPE html>
<html lang="en">

<head>
    <meta charset="UTF-8">
    <meta name="viewport" 
          content="width=device-width,
                   initial-scale=1.0">
    <title>ml5.js Custom Model Training</title>
    <style>
        #container {
            width: 320px;
            margin: auto;
            text-align: center;
            border: 1px solid #d80d0d;
        }

        canvas {
            display: block;
            margin: auto;
        }
    </style>
</head>

<body>
    <div id="container">
        <h1>Webcam Feed</h1>
        <button id="trainButton">Train Model</button>
        <button id="saveButton">Save Model</button>
    </div>
    <script src=
"https://cdnjs.cloudflare.com/ajax/libs/p5.js/1.4.0/p5.js">
    </script>
    <script src=
"https://cdn.jsdelivr.net/npm/[email protected]/dist/ml5.min.js">
    </script>
    <script src="script.js"></script>
</body>

</html>

JavaScript

let classifier;
let video;
let canvas;

function setup() {
    // Create a canvas and place it inside the container div
    canvas = createCanvas(320, 240);
    canvas.parent('container');

    // Create a video element and start the webcam
    video = createCapture(VIDEO);
    video.size(320, 240);
    // Hide the default video element
    video.hide(); 
   

    // Initialize the image classifier with MobileNet
    classifier = ml5.imageClassifier('MobileNet', video, modelReady);
}

function draw() {
    // Draw the video feed to the canvas
    image(video, 0, 0, width, height); // Ensure the video is drawn to fill the canvas
}

function modelReady() {
    console.log('Model is ready!');
}

function saveModel() {
    console.log('Model saved!');
    alert('Model saved!');
    // Placeholder for save model logic
}

function trainModel() {
    if (classifier) {
        console.log('Training the model...');
        alert('Training the model...');
        // Placeholder for training logic
    } else {
        console.error('Classifier not ready.');
    }
}

// Add event listeners after the DOM has fully loaded
window.onload = function () {
    const saveButton = document.getElementById('saveButton');
    const trainButton = document.getElementById('trainButton');

    if (saveButton && trainButton) {
        saveButton.addEventListener('click', saveModel);
        trainButton.addEventListener('click', trainModel);
    } else {
        console.error('Buttons not found in the DOM.');
    }
};

// Ensure setup is called to initialize the classifier and webcam
setup();

Note: When prompted, allow access to your webcam.

Output:

Conclusion

Training a custom model in ml5.js is a straightforward process that allows you to leverage machine learning in your web applications without extensive knowledge of the underlying algorithms. By using a pre-trained model like MobileNet and adding your own images, you can create a powerful and customized model to recognize specific objects or patterns. This process not only enhances the functionality of your application but also introduces you to the fundamentals of machine learning in a practical, hands-on way.

How to Perform Sound Classification with ml5.js?

tmishra2001

Improve

Article Tags :

How to Train a Custom Model in ml5.js?

Approach

Conclusion

Similar Reads

Thank You!

What kind of Experience do you want to share?