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 forward(self, x: ImageBatch) -> NetOutput:
fp = self.backbone(x)
fp = self.fpn(fp)
heatmaps = Heatmaps(self.hm_reg(fp[self.out_idx]))
anchors = self.anchors(heatmaps)
boxmaps = self.box_reg(fp[self.out_idx])
return heatmaps, BoxMaps(boxmaps), anchors
def __call__(self, images: ImageBatch) -> YoloBoxes:
h, w = images.shape[2:]
device = images.device
if self.use_cache:
if (h, w) in self.cache:
return self.cache[(h, w)]
grid_y, grid_x = torch.meshgrid( # type:ignore
torch.arange(h, dtype=torch.float32) / h,
torch.arange(w, dtype=torch.float32) / w,
)
box_wh = torch.tensor([1 / w, 1 / h])
box_wh = self.ratios * self.scales * box_wh
box_wh = (
box_wh.to(device)
.view(self.num_anchors, 2, 1, 1)
.expand((self.num_anchors, 2, h, w))
)
grid_xy = (
torch.stack([grid_x, grid_y]).to(device).expand(self.num_anchors, 2, h, w)
)
boxmaps = BoxMaps(torch.cat([grid_xy, box_wh], dim=1))
boxes = boxmaps_to_boxes(boxmaps)
boxes = yolo_clamp(boxes)
if self.use_cache:
self.cache[(h, w)] = boxes
return boxes
def forward(self, x: ImageBatch) -> NetOutput:
fp = self.backbone(x)
fp = self.fpn(fp)
h_fp = self.hm_reg(fp[self.out_idx])
heatmaps = self.hm_out(h_fp)
anchors = self.anchors(heatmaps)
diffmaps = self.box_out(self.box_reg(fp[self.out_idx]))
return anchors, BoxMaps(diffmaps), Heatmaps(heatmaps)
def __call__(
self,
box_batch: List[YoloBoxes],
heatmaps: Heatmaps,
) -> BoxMaps:
bms: List[torch.Tensor] = []
for boxes in box_batch:
bms.append(self._mkmaps(boxes, heatmaps))
return BoxMaps(torch.cat(bms, dim=0))
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")