A class probability map assigns an object-class prediction to each location in an image grid. In YOLO, it accompanies candidate bounding boxes and confidence scores, allowing each region to be classified. section 12.2.2 section 12.4
A YOLO prediction head
For a 6 × 6 grid and 91 COCO labels, the prediction head produces 5 box-related values plus 91 class values at every grid location. Reshaping separates the grid, and softmax converts the class values into class predictions. In R/Keras, the class output is created with `layer_activation_softmax(x@r[.., 6:NA])`. section 12.2.2
The map participates in training and visualization as follows:
- Training targets store the correct object label at grid cells intersecting each bounding box; overlapping boxes are deliberately simplified away. section 12.2.3
- The associated box target records coordinates and confidence, with confidence 1 for an actual box and 0 elsewhere. section 12.2.3
- Visualization can draw the class activation map over the image and display labels alongside predicted boxes. section 12.2.3
- Accuracy can improve through better handling of overlapping boxes or by predicting multiple boxes at each grid location. section 12.2.4
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)
Inspect the complete R/Keras construction of the YOLO prediction head, including reshaping, splitting box and class outputs, and applying softmax. section 12.2.2
