Method for object detection and recognition based on neural network
The present disclosure provides a method for object detection and recognition based on a neural network. The method includes: adding a detection layer following three detection layers of an existing YOLOv5 network model, to construct a new YOLOv5 network model; then, training the new YOLOv5 network model by considering an overlapping area between a predicted box and a ground truth box, a center-to-center distance between the two boxes, and an aspect ratio of the two boxes; and finally, inputting a to-be-detected image into the trained new YOLOv5 network model, outputting a predicted box of an object and probability values corresponding to a class to which the object belongs, and setting a class corresponding to a maximum probability value as a predicted class of the object in the to-be-detected image. This method can quickly and effectively detect multiple classes of objects. Especially, a detection effect for small objects is more ideal.
1. A method for object detection and recognition based on a neural network, comprising:
S 100 : constructing a new YOLOv5 network model by adding a detection layer following three detection layers of an existing YOLOv5 network model;
S 200 : training the new YOLOv5 network model, wherein a specific training process comprises:
S 210 : constructing a training data set: acquiring N images, resizing each of the N images to make it suitable for model training, and labeling each of the N images with ground truth boxes and object class labels, wherein all of the N labeled images constitute the training data set;
S 220 : setting thresholds for a center-to-center distance and an aspect ratio of the new YOLOv5 network model;
S 230 : initializing parameters in the new YOLOv5 network model;
inputting all samples of the training data set into the new YOLOv5 network model, performing calculation through a following formula:
s
i
=
{
s
i
,
IoU
-
R
CIoU
(
M
,
B
i
)
<
ɛ
0
,
IoU
-
R
CIoU
(
M
,
B
i
)
≥
ɛ
(
1
)
IoU represents an aspect ratio of a predicted box and a ground truth box, and is expressed by:
IoU
=
b
⋂
b
gt
b
⋃
b
gt
(
2
)
R CIoU represents a distance between a center point of the ground truth box and a center point of an overlapping area between the predicted box and the ground truth box, and is expressed by:
R
CIoU
=
ρ
2
(
b
,
b
gt
)
c
2
+
αμ
(
3
)
α
=
μ
(
1
-
IoU
)
+
μ
(
4
)
μ
=
4
π
2
(
arctan
ω
gt
h
gt
-
arctan
ω
h
)
2
(
5
)
where s i represents a classification score of an object of each class, ε represents an artificially set NMS threshold, M represents a value of the predicted box with a highest score, B i represents a list of the predicted boxes, b represents the predicted box, b gt represents the ground truth box, ρ 2 (b, b gt ) represents a distance between a center point of the predicted box and the center point of the ground truth box, and is expressed by a diagonal length of a smallest enclosing rectangular box covering the predicted box and the ground truth box, ω gt and h gt respectively indicate width and height of the ground truth box, ω and h respectively indicate width and height of the predicted box;
S 240 : during the training in S 230 , due to suppression problems of the predicted box, when a difference between IoU-CIoU of the predicted box M with the highest score and IoU-CIoU of another box B i is less than the set threshold ε, keeping a score s i of the box B i unchanged; otherwise, directly changing s i to 0, to filter out the predicted box; and
calculating loss functions, the loss functions include an object loss function, a class loss function, and a box loss function, performing repeated iterative training to minimize the loss functions, and obtaining optimal parameters of the network model; and
S 300 : detecting a to-be-detected image: resizing the to-be-detected image through the method in S 210 and inputting resized images into the trained new YOLOv5 network model for prediction, outputting the predicted box of an object and probability values corresponding to a class to which the object belongs, and setting a class corresponding to a maximum probability value as a predicted class of the object in the to-be-detected image.
2. The method for object detection and recognition based on a neural network according to claim 1 , wherein when constructing the training data set in the S 210 , Mosaic data augmentation is performed on the N labeled images to increase a data amount of the training data set.