YOLO (You Only Look Once) is a single-stage object detector: it predicts where objects are and what they are in one network evaluation. Unlike a two-stage detector, it does not first propose candidate regions and then classify them. This makes YOLO attractive for real-time detection, though speed can come with a small accuracy tradeoff. section 12.1.2 section 12.2.2
A simplified, original-style YOLO model divides a 448 × 448 image into a 6 × 6 grid. At each cell, it predicts one box, a confidence score, and class probabilities. Confidence should be near zero where no object is present. The original design predicted multiple boxes per cell; using one keeps this implementation manageable. section 12.2.2
Listing 12.6 Attaching a YOLO prediction head
grid_size <- 6L
num_labels <- 91L
inputs <- keras_input(shape = c(image_size, 3))
x <- inputs |>
backbone() |>
layer_conv_2d(512, c(3, 3), strides = c(2, 2)) |> #1
layer_flatten() |>
layer_dense(2048, activation = "relu",
kernel_initializer = "glorot_normal") |> #2
layer_dropout(0.5) |>
layer_dense(grid_size * grid_size * (num_labels + 5)) |>
layer_reshape(c(grid_size, grid_size, num_labels + 5)) #3
box_predictions <- x@r[.., 1:5] #4
class_predictions <- layer_activation_softmax(x@r[.., 6:NA])
outputs <- list(box = box_predictions, class = class_predictions)
model <- keras_model(inputs, outputs)
The R/Keras prediction head attaches to a pretrained ResNet backbone and reshapes its output into grid cells, each with five box values and 91 class values. section 12.2.2
To train the detector on COCO images:
- Convert annotated boxes into grid coordinates. Mark cells intersecting an object with its class; give cells with a target box confidence 1 and empty cells confidence 0. This simplified target preparation does not resolve overlapping boxes. section 12.2.3
- Stream and preprocess images in batches with tf.data, keeping a validation split. Batching avoids loading every decoded image into memory at once. section 12.2.3
- Train the class output with sparse categorical crossentropy and the box output with a custom loss. Apply box-position errors only where an object is present; train confidence toward zero for empty cells and toward intersection over union (IoU) for occupied cells. section 12.2.4
Listing 12.10 Defining the YOLO bounding box loss
signed_sqrt <- function(x) {
op_sign(x) * op_sqrt(op_abs(x) + config_epsilon())
}
box_loss <- function(true, pred) {
unpack <- \(x) list(x[.., 1:2], x[.., 3:4], x[.., 5:NA]) #1
.[xy_true, wh_true, conf_true] <- unpack(true) #1
.[xy_pred, wh_pred, conf_pred] <- unpack(pred) #1
no_object <- conf_true == 0 #2
xy_error <- op_square(xy_true - xy_pred) #3
wh_error <- op_square(signed_sqrt(wh_true) - signed_sqrt(wh_pred))
iou <- intersection_over_union(true, pred) #4
conf_target <- op_where(no_object, 0, op_expand_dims(iou, -1)) #4
conf_error <- op_square(conf_target - conf_pred) #4
error <- op_concatenate(axis = -1, list( #5
op_where(no_object, 0, xy_error * 5), #5
op_where(no_object, 0, wh_error * 5), #5
op_where(no_object, conf_error * 0.5, conf_error) #5
)) #5
op_sum(error, axis = c(2, 3, 4)) #6
}
The custom box loss increases the weight of box placement, reduces the influence of numerous empty cells, and compares square-root-transformed widths and heights so large boxes do not dominate. section 12.2.4
At prediction time, a confidence cutoff can hide low-confidence boxes. Viewing predictions both with and without the cutoff helps reveal how many candidate boxes the grid produces. section 12.2.4
