To train a custom YOLO (You Only Look Once) object detector in MATLAB,

by Posted on

To train a custom YOLO (You Only Look Once) object detector in MATLAB, follow these steps:

1. Prepare Your Dataset

  • Collect Images: Gather a dataset of images containing the objects you want to detect.
  • Label Data: Use the Image Labeler app (imageLabeler) to annotate objects with bounding boxes.
    • Export labels to a groundTruth object or a table with columns: imageFilename, objectBoundingBoxes.
  • Organize Data: Split into training and validation sets (e.g., 80% training, 20% validation).
matlab
Copy
% Load labeled data
data = load('labeledData.mat');
trainingData = data.gTruth.LabelData;

2. Preprocess Data

  • Resize Images: YOLO requires fixed input dimensions (e.g., 416×416).
  • Convert Labels: Convert bounding boxes to YOLO format [x_center, y_center, width, height] (normalized to image size).
  • Data Augmentation: Use transform and augmentedImageDatastore to apply augmentations (flips, rotations, etc.).
matlab
Copy
% Create an augmentedImageDatastore
augmentedTrainingData = augmentedImageDatastore([416 416], trainingData, ...
    'ColorPreprocessing', 'gray2rgb');

3. Define YOLO Architecture

  • Use yolov2ObjectDetector with a pretrained backbone (e.g., ResNet-50 or DarkNet-19).
  • Estimate Anchor Boxes: Use estimateAnchorBoxes to determine optimal box sizes for your dataset.
matlab
Copy
% Estimate anchor boxes
anchorBoxes = estimateAnchorBoxes(trainingData, 5); % 5 boxes

% Create YOLOv2 network
inputSize = [416 416 3];
numClasses = 2; % e.g., 'cat' and 'dog'
lgraph = yolov2Layers(inputSize, numClasses, anchorBoxes, 'resnet50');

4. Configure Training Options

  • Use trainingOptions to set hyperparameters (learning rate, epochs, batch size, etc.).
  • Enable GPU acceleration if available ('executionEnvironment', 'auto').
matlab
Copy
options = trainingOptions('sgdm', ...
    'InitialLearnRate', 1e-3, ...
    'MiniBatchSize', 8, ...
    'MaxEpochs', 30, ...
    'ValidationData', validationData, ...
    'Shuffle', 'every-epoch', ...
    'VerboseFrequency', 10, ...
    'ExecutionEnvironment', 'auto');

5. Train the Model

  • Use trainYOLOv2ObjectDetector to train the network.
matlab
Copy
[detector, info] = trainYOLOv2ObjectDetector(augmentedTrainingData, lgraph, options);

6. Evaluate the Model

  • Test the detector on validation images and compute metrics (precision, recall, mAP).
  • Visualize predictions with detect and insertObjectAnnotation.
matlab
Copy
% Test on a sample image
img = imread('test.jpg');
[bboxes, scores] = detect(detector, img);
annotatedImg = insertObjectAnnotation(img, 'rectangle', bboxes, scores);
imshow(annotatedImg);

7. Save/Export the Model

  • Save the trained detector for later use or deploy it.
matlab
Copy
save('customYOLOv2Detector.mat', 'detector');

Key Considerations

  • Anchor Boxes: Use estimateAnchorBoxes to optimize box sizes for your dataset.
  • Transfer Learning: Start with a pretrained backbone for faster convergence.
  • Data Augmentation: Critical for small datasets to prevent overfitting.
  • Hardware: Training is faster on GPUs (enable with Parallel Computing Toolbox).

Example Full Workflow

matlab
Copy
% Load and preprocess data
data = load('labeledData.mat');
trainingData = data.gTruth.LabelData;
trainingData.imageFilename = fullfile(data.gTruth.DataSource.Source, trainingData.imageFilename);

% Split data into training/validation
rng(0);
shuffledIdx = randperm(height(trainingData));
trainIdx = shuffledIdx(1:floor(0.8 * end));
validationIdx = shuffledIdx(floor(0.8 * end)+1:end);
trainData = trainingData(trainIdx, :);
validationData = trainingData(validationIdx, :);

% Create YOLOv2 detector
anchorBoxes = estimateAnchorBoxes(trainData, 5);
inputSize = [416 416 3];
numClasses = 2;
lgraph = yolov2Layers(inputSize, numClasses, anchorBoxes, 'resnet50');

% Train
options = trainingOptions('sgdm', 'InitialLearnRate', 1e-3, 'MaxEpochs', 30, ...);
[detector, info] = trainYOLOv2ObjectDetector(trainData, lgraph, options);

% Evaluate
results = detect(detector, validationData);
[ap, recall, precision] = evaluateDetectionPrecision(results, validationData);

Troubleshooting

  • Low Accuracy: Increase dataset size, adjust anchor boxes, or tune hyperparameters.
  • Slow Training: Use a smaller backbone (e.g., Tiny YOLO) or enable GPU.

For YOLOv3/v4 support, check MATLAB’s latest toolboxes or third-party implementations like Deep Learning Toolbox™ Converter for ONNX Model Format.

Published by howdy