| """ |
| Creates a Pytorch dataset to load the Pascal VOC & MS COCO datasets |
| """ |
|
|
| import os |
|
|
| import numpy as np |
| import pandas as pd |
| import torch |
| from PIL import Image, ImageFile |
| from torch.utils.data import Dataset |
|
|
| ImageFile.LOAD_TRUNCATED_IMAGES = True |
|
|
|
|
| class YOLODataset(Dataset): |
| def __init__( |
| self, |
| csv_file, |
| img_dir, |
| label_dir, |
| anchors, |
| image_size=416, |
| S=[13, 26, 52], |
| transform=None, |
| load_mosaic=True, |
| ): |
| self.annotations = pd.read_csv(csv_file) |
| self.img_dir = img_dir |
| self.label_dir = label_dir |
| self.image_size = image_size |
| self.transform = transform |
| self.S = S |
| self.load_mosaic = load_mosaic |
|
|
| |
| self.anchors = torch.tensor(anchors) |
| self.num_anchors_per_scale = self.anchors.shape[1] |
|
|
| def __len__(self): |
| return len(self.annotations) |
|
|
| @staticmethod |
| def iou(box, anchors): |
| """ |
| box: |
| tensor shape: [2] |
| anchors: |
| tensor shape: [number of states, number of anchors, 2] |
| |
| * 2 above is for width and height |
| """ |
|
|
| intersection = torch.prod(torch.min(box, anchors), dim=-1) |
| union = torch.prod(box) + torch.prod(anchors, dim=-1) - intersection |
| return intersection / union |
|
|
| def __getitem__(self, index): |
| label_path = os.path.join(self.label_dir, self.annotations.iloc[index, 1]) |
|
|
| |
| bboxes = np.roll( |
| np.loadtxt(fname=label_path, delimiter=" ", ndmin=2), 4, axis=1 |
| ) |
| img_path = os.path.join(self.img_dir, self.annotations.iloc[index, 0]) |
| image = np.array(Image.open(img_path).convert("RGB")) |
|
|
| if self.transform: |
| augmentations = self.transform(image=image, bboxes=bboxes) |
| image = augmentations["image"] |
| bboxes = augmentations["bboxes"] |
|
|
| """ |
| Below assumes 3 scale predictions (as paper) and same num of anchors per scale |
| 6 = [objectness, cx, cy, w, h, class] |
| """ |
| targets = [torch.zeros((self.num_anchors_per_scale, S, S, 6)) for S in self.S] |
|
|
| for bbox in bboxes: |
| iou = self.iou(torch.tensor(bbox[2:4]), self.anchors) |
|
|
| idx = torch.argsort(iou, descending=True, dim=-1) |
| idx = idx[:, 0].tolist() |
|
|
| dimensions, class_ = np.array(bbox[:-1]), bbox[-1] |
|
|
| for scale_idx, anchor_id in enumerate(idx): |
| scale_dim = self.S[scale_idx] |
| scale_cx, scale_cy, scale_w, scale_h = dimensions * scale_dim |
|
|
| row, col = int(scale_cy), int(scale_cx) |
|
|
| |
| scale_cx = scale_cx - col |
| scale_cy = scale_cy - row |
|
|
| box_target = torch.tensor( |
| [1, scale_cx, scale_cy, scale_w, scale_h, class_] |
| ) |
|
|
| targets[scale_idx][anchor_id, row, col] = box_target |
|
|
| return image, targets |
|
|
|
|
| if __name__ == "__main__": |
| from src.run.yolov3 import config |
|
|
| IMAGE_SIZE = config.IMAGE_SIZE |
| train_dataset = YOLODataset( |
| config.DATASET + "/2examples.csv", |
| transform=config.train_transforms, |
| S=[IMAGE_SIZE // 32, IMAGE_SIZE // 16, IMAGE_SIZE // 8], |
| img_dir=config.IMG_DIR, |
| label_dir=config.LABEL_DIR, |
| anchors=config.ANCHORS, |
| ) |
|
|
