def build_infer_batch(records):
imgs, img_metas = [], []
for record in records:
imgs.append(_img_tensor(record))
img_metas.append(_img_meta_mask(record))
data = {
"img": [torch.stack(imgs)],
"img_metas": [img_metas],
}
return data, records
def build_infer_batch(records, batch_tfms=None):
records = common_build_batch(records, batch_tfms=batch_tfms)
imgs, img_metas = [], []
for record in records:
imgs.append(_img_tensor(record))
img_metas.append(_img_meta_mask(record))
data = {
"img": [torch.stack(imgs)],
"img_metas": [img_metas],
}
return data, records
def build_train_batch(
records: Sequence[RecordType],
) -> Tuple[dict, List[Dict[str, torch.Tensor]]]:
images, labels, bboxes, masks, img_metas = [], [], [], [], []
for record in records:
images.append(_img_tensor(record))
img_metas.append(_img_meta_mask(record))
labels.append(_labels(record))
bboxes.append(_bboxes(record))
masks.append(_masks(record))
data = {
"img": torch.stack(images),
"img_metas": img_metas,
"gt_labels": labels,
"gt_bboxes": bboxes,
"gt_masks": masks,
}
return data, records