class AssignGTtoDefaultBox(object): def __init__(self): self.default_box = dboxes300_coco() self.encoder = Encoder(self.default_box) def __call__(self, image, target): boxes = target['boxes'] labels = target["labels"] # bboxes_out (Tensor 8732 x 4), labels_out (Tensor 8732) bboxes_out, labels_out = self.encoder.encode(boxes, labels) target['boxes'] = bboxes_out target['labels'] = labels_out return image, target
class SSDTransformer(object): def __init__(self, dboxes, size=(300, 300), val=False): self.size = size self.val = val self.dboxes = dboxes self.encoder = Encoder(self.dboxes) self.crop = SSDCropping() self.hflip = RandomHorizontalFlip() self.normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) self.img_trans = transforms.Compose([ transforms.Resize(self.size), transforms.ColorJitter(brightness=0.125, contrast=0.5, saturation=0.5, hue=0.05), transforms.ToTensor(), self.normalize ]) self.trans_val = transforms.Compose([ transforms.Resize(self.size), transforms.ToTensor(), self.normalize ]) def __call__(self, img, img_size, bboxes=None, labels=None, max_num=200): if self.val: bbox_out = torch.zeros(max_num, 4) label_out = torch.zeros(max_num, dtype=torch.long) bbox_out[:bboxes.size(0), :] = bboxes label_out[:labels.size(0)] = labels return self.trans_val(img), img_size, bbox_out, label_out img, img_size, bboxes, labels = self.crop(img, img_size, bboxes, labels) img, bboxes = self.hflip(img, bboxes) img = self.img_trans(img).contiguous() bboxes, labels = self.encoder.encode(bboxes, labels) return img, img_size, bboxes, labels