def get_targets(self, target, sample):
batched_preds = []
batched_offsets = []
for t, s in zip(target, sample):
bboxes = t['bounding_box'].to(self.device)
gt_classes, gt_offsets = get_bbox_gt(bboxes, t['category'].to(self.device), self.anchor_boxes.to(self.device), self.map_sz, self.device)
batched_preds.append(gt_classes)
batched_offsets.append(gt_offsets)
class_targets = torch.stack(batched_preds)
box_targets = torch.stack(batched_offsets)
return class_targets, box_targets
def my_collate_fn(batch):
images = []
gt_boxes = get_anchor_boxes()
img_h = 256
img_w = 306
map_sz = 800
class_target = []
box_target = []
for x in batch:
images.append(x[0])
gt_classes, gt_offsets = get_bbox_gt(x[1]['bounding_box'],
x[1]['category'], gt_boxes,
map_sz)
class_target.append(gt_classes)
box_target.append(gt_offsets)
samples = torch.stack(images)
samples = samples.view(len(batch), -1, img_h, img_w).double()
class_target = torch.stack(class_target)
box_target = torch.stack(box_target)
return samples, class_target, box_target