Addition of additional CNN and Dense lyers in the existing YOLOv4 model layers for subclass level detection

Hello,

I am working on an object detection problem where it is required to detect objects at two levels. At the first level, it has classes and at the second level, it has subclasses. For example, classes are military ships or commercial ships or private ships. And subclasses for military ships are e.g., Destroyer, frigates, or corvette. For subclass level detection, it is required to add extra CNN and Dense layers to the subclass level detection.

In the present YOLOv4 code, I am not able to modify or add extralayers in the neural network layer architecture.

Kindly suggest/guide for implementation.

Answer 1

[Maneet Kaur Bagga]

Hi NARENDRA,

As per my understanding, adding additional CNN layers directly to the YOLO model can disrupt the anchor box mechanism and subsequent processing steps leading to degraded performance or incorrect predictions. Alternatively, subclass level detection can be performed by a two-step approach as below:

First Level Object Detection :

• Training the model using YOLO V4 model to detect the classes at first level such as military ships, commercial ships and private ships.
• This model will provide bounding box predictions and confidence scores for each detected object at the first level.

Second Level Object Detection :

• Once the bounding box predictions and confidence scores of the first level classes are obtained, crop the Regions of Interest(ROI) for each detected object.
• Resize these ROIs to a fixed size and feed them into a separate CNN model. Add additional CNN layers and Dense layers to this separate CNN model to perform subclass level detection.
• Train this separate CNN model on a dataset containing subclasses, for example destroyer, frigates, corvette within each first level class. Use this trained model to predict the subclasses within each first level object.

Please refer to the following code snippet for better understanding of the above approach:

% Step 1: First-Level Object Detection (YOLOv4)

% Load and preprocess the dataset for first-level object detection
datasetPath = 'path/to/first_level_dataset';
imds = imageDatastore(datasetPath, 'IncludeSubfolders', true, 'LabelSource', 'foldernames');
imds.ReadFcn = @readAndPreprocessImage;

% Split the dataset into training and validation sets
[imdsTrain, imdsVal] = splitEachLabel(imds, 0.8, 'randomized');

% Define the YOLOv4 network architecture
inputSize = [416 416 3];
numClasses = 3; % Number of first-level classes
lgraph = yolov4Layers(inputSize, numClasses);

% Set training options
options = trainingOptions('adam', ...
    'InitialLearnRate', 1e-4, ...
    'MiniBatchSize', 8, ...
    'MaxEpochs', 50, ...
    'Shuffle', 'every-epoch', ...
    'Verbose', true, ...
    'Plots', 'training-progress', ...
    'ValidationData', imdsVal, ...
    'ValidationFrequency', 10);

% Train the YOLOv4 model
detector = trainYOLOv2ObjectDetector(imdsTrain, lgraph, options);

% Step 2: Second-Level Object Detection (Subclass Detection)

% Load and preprocess the dataset for subclass detection
subclassDatasetPath = 'path/to/subclass_dataset';
subclassImds = imageDatastore(subclassDatasetPath, 'IncludeSubfolders', true, 'LabelSource', 'foldernames');
subclassImds.ReadFcn = @readAndPreprocessImage;

% For each first-level class, train a separate subclass detection model
firstLevelClasses = {'military_ships', 'commercial_ships', 'private_ships'};
numSubclasses = [3, 4, 2]; % Number of subclasses for each first-level class

for i = 1:numel(firstLevelClasses)
    % Filter the subclass dataset for the current first-level class
    currentClass = firstLevelClasses{i};
    currentSubclassImds = subset(subclassImds, @(labels) contains(labels, currentClass));
    
    % Prepare the dataset for subclass detection
    subclassData = objectDetectorTrainingData(currentSubclassImds, 'DataAugmentation', 'random');
    
    % Define the subclass detection network architecture
    inputSize = [224 224 3];
    numClasses = numSubclasses(i);
    subclassLayers = [
        imageInputLayer(inputSize)
        convolution2dLayer(3, 64, 'Padding', 'same')
        reluLayer
        maxPooling2dLayer(2, 'Stride', 2)
        % Add more layers as needed for subclass detection
        fullyConnectedLayer(numClasses)
        softmaxLayer
        classificationLayer];
    
    % Set training options for subclass detection
    subclassOptions = trainingOptions('adam', ...
        'InitialLearnRate', 1e-4, ...
        'MiniBatchSize', 32, ...
        'MaxEpochs', 20, ...
        'Verbose', true, ...
        'Plots', 'training-progress');
    
    % Train the subclass detection model
    subclassDetector = trainNetwork(subclassData, subclassLayers, subclassOptions);
    
    % Save the trained subclass detection model
    save(fullfile('path/to/save/models', [currentClass '_subclass_detector.mat']), 'subclassDetector');
end

Refer to the below file-exchange link to implement YOLO V4 model using MATALB

https://www.mathworks.com/matlabcentral/fileexchange/75305-yolov3-yolov4-matlab?s_tid=answers_rc2-2_p5_MLT

Hope this helps!

Thank You

Maneet Bagga