| 54 | |
| 55 | |
| 56 | class VOCDataset(Dataset): |
| 57 | def __init__(self, split, im_dir, ann_dir): |
| 58 | self.split = split |
| 59 | self.im_dir = im_dir |
| 60 | self.ann_dir = ann_dir |
| 61 | classes = [ |
| 62 | 'person', 'bird', 'cat', 'cow', 'dog', 'horse', 'sheep', |
| 63 | 'aeroplane', 'bicycle', 'boat', 'bus', 'car', 'motorbike', 'train', |
| 64 | 'bottle', 'chair', 'diningtable', 'pottedplant', 'sofa', 'tvmonitor' |
| 65 | ] |
| 66 | classes = sorted(classes) |
| 67 | classes = ['background'] + classes |
| 68 | self.label2idx = {classes[idx]: idx for idx in range(len(classes))} |
| 69 | self.idx2label = {idx: classes[idx] for idx in range(len(classes))} |
| 70 | print(self.idx2label) |
| 71 | self.images_info = load_images_and_anns(im_dir, ann_dir, self.label2idx) |
| 72 | |
| 73 | def __len__(self): |
| 74 | return len(self.images_info) |
| 75 | |
| 76 | def __getitem__(self, index): |
| 77 | im_info = self.images_info[index] |
| 78 | im = Image.open(im_info['filename']) |
| 79 | to_flip = False |
| 80 | if self.split == 'train' and random.random() < 0.5: |
| 81 | to_flip = True |
| 82 | im = im.transpose(Image.FLIP_LEFT_RIGHT) |
| 83 | im_tensor = torchvision.transforms.ToTensor()(im) |
| 84 | targets = {} |
| 85 | targets['bboxes'] = torch.as_tensor([detection['bbox'] for detection in im_info['detections']]) |
| 86 | targets['labels'] = torch.as_tensor([detection['label'] for detection in im_info['detections']]) |
| 87 | if to_flip: |
| 88 | for idx, box in enumerate(targets['bboxes']): |
| 89 | x1, y1, x2, y2 = box |
| 90 | w = x2-x1 |
| 91 | im_w = im_tensor.shape[-1] |
| 92 | x1 = im_w - x1 - w |
| 93 | x2 = x1 + w |
| 94 | targets['bboxes'][idx] = torch.as_tensor([x1, y1, x2, y2]) |
| 95 | return im_tensor, targets, im_info['filename'] |
| 96 |
no outgoing calls
no test coverage detected