def __call__(
self, preds: BoxMaps, gt_box_batch: List[YoloBoxes], anchormap: BoxMap
) -> Tensor:
device = preds.device
_, _, h, w = preds.shape
box_losses: List[Tensor] = []
anchors = boxmap_to_boxes(anchormap)
for diff_map, gt_boxes in zip(preds, gt_box_batch):
if len(gt_boxes) == 0:
continue
pred_boxes = boxmap_to_boxes(BoxMap(diff_map))
match_indices, positive_indices = self.matcher(anchors, gt_boxes, (w, h))
num_pos = positive_indices.sum()
if num_pos == 0:
continue
matched_gt_boxes = YoloBoxes(gt_boxes[match_indices][positive_indices])
matched_pred_boxes = YoloBoxes(pred_boxes[positive_indices])
if self.use_diff:
matched_pred_boxes = YoloBoxes(
anchors[positive_indices] + matched_pred_boxes
)
box_losses.append(
self.loss(
yolo_to_pascal(matched_pred_boxes, (1, 1)),
yolo_to_pascal(matched_gt_boxes, (1, 1)),
)
)
if len(box_losses) == 0:
return torch.tensor(0.0).to(device)
return torch.stack(box_losses).mean()
def test_nearest_assign() -> None:
x = YoloBoxes(
torch.tensor(
[
[0.11, 0.12, 0.1, 0.1],
[0.21, 0.22, 0.1, 0.1],
[0.25, 0.22, 0.1, 0.1],
[0.31, 0.32, 0.1, 0.1],
[0.41, 0.42, 0.1, 0.1],
[0.61, 0.62, 0.1, 0.1],
]
)
)
y = YoloBoxes(
torch.tensor(
[[0.1, 0.1, 0.1, 0.1], [0.2, 0.2, 0.1, 0.1], [0.3, 0.3, 0.1, 0.1],]
)
)
fn = NearnestAssign()
matched_idx, positive_idx = fn(x, y)
assert (
F.l1_loss(matched_idx.float(), torch.tensor([0, 1, 1, 2, 2, 2]).float()) < 1e-9
)
assert (
F.l1_loss(positive_idx.float(), torch.tensor([1, 1, 0, 1, 0, 0]).float()) < 1e-9
)
def __call__(
self,
inputs: NetOutput) -> Tuple[List[YoloBoxes], List[Confidences]]:
anchormap, box_diffs, heatmap = inputs
device = heatmap.device
if self.use_peak:
kpmap = (self.max_pool(heatmap)
== heatmap) & (heatmap > self.threshold)
else:
kpmap = heatmap > self.threshold
batch_size, _, height, width = heatmap.shape
original_wh = torch.tensor([width, height],
dtype=torch.float32).to(device)
rows: List[Tuple[YoloBoxes, Confidences]] = []
box_batch = []
conf_batch = []
for hm, km, box_diff in zip(heatmap.squeeze(1), kpmap.squeeze(1),
box_diffs):
kp = torch.nonzero(km, as_tuple=False)
confidences = hm[kp[:, 0], kp[:, 1]]
anchor = anchormap[:, kp[:, 0], kp[:, 1]].t()
box_diff = box_diff[:, kp[:, 0], kp[:, 1]].t()
boxes = anchor + box_diff
boxes = yolo_clamp(YoloBoxes(boxes))
sort_idx = nms(yolo_to_pascal(boxes, (1, 1)), confidences,
self.nms_thresold)
box_batch.append(YoloBoxes(boxes[sort_idx]))
conf_batch.append(Confidences(confidences[sort_idx]))
return box_batch, conf_batch
def __call__(
self, net_output: NetOutput
) -> t.Tuple[List[YoloBoxes], List[Confidences]]:
anchor_levels, box_diff_levels, labels_levels = net_output
box_batch = []
confidence_batch = []
anchors = torch.cat(anchor_levels, dim=0) # type: ignore
box_diffs = torch.cat(box_diff_levels, dim=1) # type:ignore
labels_batch = torch.cat(labels_levels, dim=1) # type:ignore
for box_diff, confidences in zip(box_diffs, labels_batch):
boxes = anchors + box_diff
confidences, c_index = confidences.max(dim=1)
filter_idx = confidences > self.confidence_threshold
confidences = confidences[filter_idx][:self.limit]
boxes = boxes[filter_idx][:self.limit]
sort_idx = nms(
yolo_to_pascal(YoloBoxes(boxes), (1, 1)),
confidences,
self.iou_threshold,
)
confidences.argsort(descending=True)
boxes = YoloBoxes(boxes[sort_idx])
boxes = yolo_clamp(boxes)
confidences = confidences[sort_idx]
box_batch.append(boxes)
confidence_batch.append(Confidences(confidences))
return box_batch, confidence_batch
def __getitem__(self, idx: int) -> TrainSample:
image = torch.rand((3, *self.image_size), dtype=torch.float32)
boxes = torch.rand((random.randint(1, 9), 4),
dtype=torch.float32).clamp(0, 1.0)
labels = torch.zeros((len(boxes), ))
return (ImageId(""), Image(image), YoloBoxes(boxes.float()),
Labels(labels))
def __call__(
self, batch: t.Tuple[ImageBatch, t.List[YoloBoxes]]
) -> t.Tuple[ImageBatch, List[YoloBoxes]]:
image_batch, box_batch = batch
image_batch = ImageBatch(image_batch.to(self.device))
box_batch = [YoloBoxes(x.to(self.device)) for x in box_batch]
return image_batch, box_batch
def test_anchors() -> None:
fn = Anchors(size=2)
hm = torch.zeros((1, 1, 8, 8))
anchors = fn(hm)
boxes = boxmap_to_boxes(anchors)
assert len(boxes) == 8 * 8
plot = DetectionPlot(w=8, h=8)
plot.with_yolo_boxes(YoloBoxes(boxes[[0, 4, 28, 27]]), color="blue")
plot.save(f"store/test-anchorv1.png")
def test_mkcornermaps(h: int, w: int, cy: int, cx: int, dy: float, dx: float) -> None:
in_boxes = YoloBoxes(torch.tensor([[0.201, 0.402, 0.1, 0.3]]))
to_boxes = ToBoxes(threshold=0.1)
mkmaps = MkCornerMaps()
hm = mkmaps([in_boxes], (h, w), (h * 10, w * 10))
assert hm.shape == (1, 1, h, w)
mk_anchors = Anchors()
anchormap = mk_anchors(hm)
diffmaps = BoxMaps(torch.zeros((1, *anchormap.shape)))
diffmaps = in_boxes.view(1, 4, 1, 1).expand_as(diffmaps) - anchormap
out_box_batch, out_conf_batch = to_boxes((anchormap, diffmaps, hm))
out_boxes = out_box_batch[0]
for box in out_boxes:
assert F.l1_loss(box, in_boxes[0]) < 1e-8
plot = DetectionPlot(w=w, h=h)
plot.with_image((hm[0, 0] + 1e-4).log())
plot.with_yolo_boxes(out_boxes, color="red")
plot.with_yolo_boxes(in_boxes, color="blue")
plot.save(f"store/test-corner.png")
def test_mkmaps(h: int, w: int, cy: int, cx: int, dy: float, dx: float) -> None:
in_boxes = YoloBoxes(torch.tensor([[0.201, 0.402, 0.1, 0.3]]))
to_boxes = ToBoxes(threshold=0.1)
mkmaps = MkGaussianMaps(sigma=2.0)
hm = mkmaps([in_boxes], (h, w), (h * 10, w * 10))
assert (torch.nonzero(hm.eq(1), as_tuple=False)[0, 2:] - torch.tensor([[cy, cx]])).sum() == 0 # type: ignore
assert hm.shape == (1, 1, h, w)
mk_anchors = Anchors()
anchormap = mk_anchors(hm)
diffmaps = BoxMaps(torch.zeros((1, *anchormap.shape)))
diffmaps = in_boxes.view(1, 4, 1, 1).expand_as(diffmaps) - anchormap
out_box_batch, out_conf_batch = to_boxes((anchormap, diffmaps, hm))
out_boxes = out_box_batch[0]
for box in out_boxes:
assert F.l1_loss(box, in_boxes[0]) < 1e-8
plot = DetectionPlot(w=w, h=h)
plot.with_image((hm[0, 0] + 1e-4).log())
plot.with_yolo_boxes(in_boxes, color="blue")
plot.with_yolo_boxes(out_boxes, color="red")
plot.save(f"store/test-heatmapv1.png")
def _mkmaps(self, boxes: YoloBoxes, heatmaps: Heatmaps) -> BoxMaps:
device = boxes.device
_, _, h, w = heatmaps.shape
boxmaps = torch.zeros((1, 4, h, w), dtype=torch.float32).to(device)
box_count = len(boxes)
if box_count == 0:
return BoxMaps(boxmaps)
wh = torch.tensor([w, h]).to(device)
cxcy = (boxes[:, :2] * wh).long()
cx = cxcy[:, 0]
cy = cxcy[:, 1]
boxmaps[:, :, cy, cx] = boxes.t()
return BoxMaps(boxmaps)
def __call__(
self, batch: t.Tuple[ImageBatch, List[YoloBoxes], List[Labels]]
) -> t.Tuple[ImageBatch, List[YoloBoxes], List[Labels]]:
image_batch, boxes_batch, label_batch = batch
return (
ImageBatch(
image_batch.to(self.device, non_blocking=self.non_blocking)),
[
YoloBoxes(x.to(self.device, non_blocking=self.non_blocking))
for x in boxes_batch
],
[
Labels(x.to(self.device, non_blocking=self.non_blocking))
for x in label_batch
],
)
def test_anchors(fsize: int, size: float, scales: List[float],
ratios: List[float]) -> None:
h = fsize
w = fsize
images = ImageBatch(torch.zeros((1, 3, h, w), dtype=torch.float32))
fn = Anchors(size=size, scales=scales, ratios=ratios)
res = fn(images)
num_anchors = len(scales) * len(ratios)
anchor_count = w * h * num_anchors
assert res.shape == (anchor_count, 4)
offset = num_anchors * h * (w // 2) + num_anchors * h // 2
ids = [offset + x for x in range(num_anchors)]
plot = DetectionPlot(w=256, h=256)
plot.with_yolo_boxes(YoloBoxes(res[ids]), color="red")
plot.save(
f"store/test-anchors-{fsize}-{size}-{'-'.join([str(x) for x in scales])}-{num_anchors}.png"
)
def __call__(self, inputs: NetOutput) -> t.List[t.Tuple[YoloBoxes, Confidences]]:
heatmaps, boxmaps, anchormap = inputs
device = heatmaps.device
kpmap = (self.max_pool(heatmaps) == heatmaps) & (heatmaps > self.threshold)
batch_size, _, height, width = heatmaps.shape
original_wh = torch.tensor([width, height], dtype=torch.float32).to(device)
rows: t.List[t.Tuple[YoloBoxes, Confidences]] = []
for hm, km, bm in zip(heatmaps.squeeze(1), kpmap.squeeze(1), boxmaps):
kp = torch.nonzero(km, as_tuple=False)
confidences = hm[kp[:, 0], kp[:, 1]]
if self.use_diff:
boxes = (
anchormap[:, kp[:, 0], kp[:, 1]].t() + bm[:, kp[:, 0], kp[:, 1]].t()
)
else:
boxes = bm[:, kp[:, 0], kp[:, 1]].t()
sort_idx = confidences.argsort(descending=True)[: self.limit]
rows.append(
(YoloBoxes(boxes[sort_idx]), Confidences(confidences[sort_idx]))
)
return rows