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 _collate_fn(batch: List[TrainSample],) -> Tuple[ImageBatch, List[ImageId]]:
images: List[Any] = []
id_batch: List[ImageId] = []
for id, img, _, _ in batch:
images.append(img)
id_batch.append(id)
return ImageBatch(torch.stack(images)), id_batch
def __call__(
self, model: nn.Module,
images: ImageBatch) -> Tuple[List[YoloBoxes], List[Confidences]]:
images = ImageBatch(self.img_transform(images))
outputs = model(images)
box_batch, conf_batch = self.to_boxes(outputs)
box_batch = [self.box_transform(boxes) for boxes in box_batch]
return box_batch, conf_batch
def test_hfliptta() -> None:
to_boxes = ToBoxes(threshold=0.1)
fn = HFlipTTA(to_boxes)
images = ImageBatch(torch.zeros((1, 3, 64, 64)))
images[:, :, 0, 0] = torch.ones((1, 3))
channels = 32
backbone = ResNetBackbone("resnet34", out_channels=channels)
model = CenterNetV1(channels=channels, backbone=backbone, out_idx=6,)
fn(model, images)
def collate_fn(
batch: List[TrainSample],
) -> Tuple[List[ImageId], ImageBatch, List[YoloBoxes], List[Labels]]:
images: List[t.Any] = []
id_batch: List[ImageId] = []
box_batch: List[YoloBoxes] = []
label_batch: List[Labels] = []
for id, img, boxes, labels in batch:
images.append(img)
box_batch.append(boxes)
id_batch.append(id)
label_batch.append(labels)
return id_batch, ImageBatch(torch.stack(images)), box_batch, label_batch
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,
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")