%matplotlib inline
from d2l import jax as d2l
import jax
from jax import numpy as jnp
from flax import linen as nn
import optax
import numpy as np
np.set_printoptions(2) # Simplify printing accuracyCoordinates and feature maps
Anchors at one pixel
Visualize the n + m - 1 anchors centered at a single pixel — different scales and ratios:
def show_bboxes(axes, bboxes, labels=None, colors=None):
"""Show bounding boxes."""
def make_list(obj, default_values=None):
if obj is None:
obj = default_values
elif not isinstance(obj, (list, tuple)):
obj = [obj]
return obj
labels = make_list(labels)
colors = make_list(colors, ['b', 'g', 'r', 'm', 'c'])
for i, bbox in enumerate(bboxes):
color = colors[i % len(colors)]
rect = d2l.bbox_to_rect(d2l.numpy(bbox), color)
axes.add_patch(rect)
if labels and len(labels) > i:
text_color = 'k' if color == 'w' else 'w'
axes.text(rect.xy[0], rect.xy[1], labels[i],
va='center', ha='center', fontsize=9, color=text_color,
bbox=dict(facecolor=color, lw=0))d2l.set_figsize()
bbox_scale = d2l.tensor((w, h, w, h))
fig = d2l.plt.imshow(img)
show_bboxes(fig.axes, boxes[250, 250, :, :] * bbox_scale,
['s=0.75, r=1', 's=0.5, r=1', 's=0.25, r=1', 's=0.75, r=2',
's=0.75, r=0.5'])
Intersection over Union
We need a similarity measure between two boxes to know which anchor matches which ground truth.
\text{IoU}(A, B) = \frac{|A \cap B|}{|A \cup B|}.
IoU = intersection area / union area.
def box_iou(boxes1, boxes2):
"""Compute pairwise IoU across two lists of anchor or bounding boxes."""
box_area = lambda boxes: ((boxes[:, 2] - boxes[:, 0]) *
(boxes[:, 3] - boxes[:, 1]))
# Shape of `boxes1`, `boxes2`, `areas1`, `areas2`: (no. of boxes1, 4),
# (no. of boxes2, 4), (no. of boxes1,), (no. of boxes2,)
areas1 = box_area(boxes1)
areas2 = box_area(boxes2)
# Shape of `inter_upperlefts`, `inter_lowerrights`, `inters`: (no. of
# boxes1, no. of boxes2, 2)
inter_upperlefts = jnp.maximum(boxes1[:, None, :2], boxes2[:, :2])
inter_lowerrights = jnp.minimum(boxes1[:, None, 2:], boxes2[:, 2:])
inters = jnp.clip(inter_lowerrights - inter_upperlefts, 0)
# Shape of `inter_areas` and `union_areas`: (no. of boxes1, no. of boxes2)
inter_areas = inters[:, :, 0] * inters[:, :, 1]
union_areas = areas1[:, None] + areas2 - inter_areas
return inter_areas / union_areasMatching anchors to ground truth
For each anchor box, decide which ground-truth box (if any) it should learn to predict. Common rule: greedy assignment by highest IoU, with a threshold (e.g. 0.5) for “positive” matches:
Anchor → GT assignment by IoU.
def assign_anchor_to_bbox(ground_truth, anchors, device, iou_threshold=0.5):
"""Assign closest ground-truth bounding boxes to anchor boxes."""
num_anchors, num_gt_boxes = anchors.shape[0], ground_truth.shape[0]
jaccard = box_iou(anchors, ground_truth)
anchors_bbox_map = jnp.full((num_anchors,), -1, dtype=jnp.int32)
max_ious = jnp.max(jaccard, axis=1)
indices = jnp.argmax(jaccard, axis=1)
mask = max_ious >= iou_threshold
anchors_bbox_map = jnp.where(mask, indices, anchors_bbox_map)
col_discard = jnp.full((num_anchors,), -1.0)
row_discard = jnp.full((num_gt_boxes,), -1.0)
# Use lax.fori_loop so JIT does not unroll (and re-trace) per gt box
def body(_, carry):
jaccard, anchors_bbox_map = carry
max_idx = jnp.argmax(jaccard)
box_idx = max_idx % num_gt_boxes
anc_idx = max_idx // num_gt_boxes
anchors_bbox_map = anchors_bbox_map.at[anc_idx].set(box_idx)
jaccard = jaccard.at[:, box_idx].set(col_discard)
jaccard = jaccard.at[anc_idx, :].set(row_discard)
return (jaccard, anchors_bbox_map)
_, anchors_bbox_map = jax.lax.fori_loop(
0, num_gt_boxes, body, (jaccard, anchors_bbox_map))
return anchors_bbox_mapWorked example
Hand-pick ground truth (dog, cat) and a few anchors; plot them:
ground_truth = d2l.tensor([[0, 0.1, 0.08, 0.52, 0.92],
[1, 0.55, 0.2, 0.9, 0.88]])
anchors = d2l.tensor([[0, 0.1, 0.2, 0.3], [0.15, 0.2, 0.4, 0.4],
[0.63, 0.05, 0.88, 0.98], [0.66, 0.45, 0.8, 0.8],
[0.57, 0.3, 0.92, 0.9]])
fig = d2l.plt.imshow(img)
show_bboxes(fig.axes, ground_truth[:, 1:] * bbox_scale, ['dog', 'cat'], 'k')
show_bboxes(fig.axes, anchors * bbox_scale, ['0', '1', '2', '3', '4']);
Run the labeler:
labels = multibox_target(jnp.expand_dims(anchors, axis=0),
jnp.expand_dims(ground_truth, axis=0))labels[2]Array([[0, 1, 2, 0, 2]], dtype=int32)NMS demo
Four overlapping predictions; NMS picks the top-scoring one and suppresses the rest. If a lower-score box covers the same object, it should disappear; if it covers a different object, its IoU should be low enough to survive:
anchors = d2l.tensor([[0.1, 0.08, 0.52, 0.92], [0.08, 0.2, 0.56, 0.95],
[0.15, 0.3, 0.62, 0.91], [0.55, 0.2, 0.9, 0.88]])
offset_preds = d2l.tensor([0] * d2l.size(anchors))
cls_probs = d2l.tensor([[0] * 4, # Predicted background likelihood
[0.9, 0.8, 0.7, 0.1], # Predicted dog likelihood
[0.1, 0.2, 0.3, 0.9]]) # Predicted cat likelihoodfig = d2l.plt.imshow(img)
show_bboxes(fig.axes, anchors * bbox_scale,
['dog=0.9', 'dog=0.8', 'dog=0.7', 'cat=0.9'])
End-to-end output
Each output row is (class_id, confidence, x1, y1, x2, y2). Rows with class -1 have been suppressed or filtered out; the remaining rows are the detector’s final boxes.
output = multibox_detection(jnp.expand_dims(cls_probs, axis=0),
jnp.expand_dims(offset_preds, axis=0),
jnp.expand_dims(anchors, axis=0),
nms_threshold=0.5)
outputArray([[[ 0. , 0.9 , 0.1 , 0.08, 0.52, 0.92],
[ 1. , 0.9 , 0.55, 0.2 , 0.9 , 0.88],
[-1. , 0.8 , 0.08, 0.2 , 0.56, 0.95],
[-1. , 0.7 , 0.15, 0.3 , 0.62, 0.91]]], dtype=float32)fig = d2l.plt.imshow(img)
for i in d2l.numpy(output[0]):
if i[0] == -1:
continue
label = ('dog=', 'cat=')[int(i[0])] + str(i[1])
show_bboxes(fig.axes, [d2l.tensor(i[2:]) * bbox_scale], label)