def __call__(
self,
net_out: NetOutput,
src: Tuple[List[YoloBoxes], List[Confidences]],
tgt: List[YoloBoxes],
image_batch: ImageBatch,
gt_hms: Heatmaps,
) -> None:
heatmap, _, _ = net_out
box_batch, confidence_batch = src
box_batch = box_batch[: self.limit]
confidence_batch = confidence_batch[: self.limit]
_, _, h, w = image_batch.shape
for i, (sb, sc, tb, hm, img, gt_hm) in enumerate(
zip(box_batch, confidence_batch, tgt, heatmap, image_batch, gt_hms)
):
plot = DetectionPlot(
h=h,
w=w,
use_alpha=self.use_alpha,
figsize=self.figsize,
show_probs=self.show_probs,
)
plot.with_image(img, alpha=0.7)
plot.with_image(hm[0].log(), alpha=0.3)
plot.with_yolo_boxes(tb, color="blue")
plot.with_yolo_boxes(sb, sc, color="red")
plot.save(f"{self.out_dir}/{self.prefix}-boxes-{i}.png")
plot = DetectionPlot(
h=h, w=w, use_alpha=self.use_alpha, figsize=self.figsize
)
plot.with_image(img, alpha=0.7)
plot.with_image((gt_hm[0] + 1e-4).log(), alpha=0.3)
plot.save(f"{self.out_dir}/{self.prefix}-hm-{i}.png")
def test_yolo_vflip() -> None:
in_boxes = YoloBoxes(torch.tensor([
[0.1, 0.1, 0.1, 0.05],
]))
out_boxes = yolo_vflip(in_boxes)
assert (out_boxes - torch.tensor([[0.1, 0.9, 0.1, 0.05]])).abs().sum() == 0
plot = DetectionPlot(w=100, h=100)
plot.with_yolo_boxes(in_boxes, color="blue")
plot.with_yolo_boxes(out_boxes, color="red")
plot.save(f"store/test-yolo-vflip.png")
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_train_dataset() -> None:
dataset = WheatDataset(config.annot_file,
config.train_image_dir,
max_size=1024)
image_id, img, boxes, _ = dataset[0]
assert img.dtype == torch.float32
assert boxes.dtype == torch.float32
for i in range(10):
_, img, boxes, _ = dataset[100]
_, h, w = img.shape
plot = DetectionPlot(figsize=(20, 20), w=w, h=h)
plot.with_image(img)
plot.with_yolo_boxes(boxes, color="red")
plot.save(f"{config.working_dir}/test-dataset-{i}.png")
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 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 __call__(
self,
src: Tuple[List[YoloBoxes], List[Confidences]],
tgt: List[YoloBoxes],
image_batch: ImageBatch,
) -> None:
image_batch = ImageBatch(image_batch[:self.limit])
tgt = tgt[:self.limit]
src_box_batch, src_confidence_batch = src
_, _, h, w = image_batch.shape
for i, (img, sb, sc, tb) in enumerate(
zip(image_batch, src_box_batch, src_confidence_batch, tgt)):
plot = DetectionPlot(h=h,
w=w,
use_alpha=self.use_alpha,
show_probs=self.show_probs)
plot.with_image(img, alpha=0.5)
plot.with_yolo_boxes(tb, color="blue")
plot.with_yolo_boxes(sb, sc, color="red")
plot.save(f"{self.out_dir}/{self.prefix}-boxes-{i}.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")