ManjunathReddy/Yolo3_from_scratch
0
1"""2Creates a Pytorch dataset to load the Pascal VOC & MS COCO datasets3"""4 5import src.config as config6import numpy as np7import os8import pandas as pd9import torch10from src.utils_rh import xywhn2xyxy, xyxy2xywhn11import random 12 13from PIL import Image, ImageFile14from torch.utils.data import Dataset, DataLoader15from src.utils_rh import (16 cells_to_bboxes,17 iou_width_height as iou,18 non_max_suppression as nms,19 plot_image20)21 22ImageFile.LOAD_TRUNCATED_IMAGES = True23 24class YOLODataset(Dataset):25 def __init__(26 self,27 csv_file,28 img_dir,29 label_dir,30 anchors,31 image_size=416,32 S=[13, 26, 52],33 C=20,34 transform=None,35 ):36 self.annotations = pd.read_csv(csv_file)37 self.img_dir = img_dir38 self.label_dir = label_dir39 self.image_size = image_size40 self.mosaic_border = [image_size // 2, image_size // 2]41 self.transform = transform42 self.S = S43 self.anchors = torch.tensor(anchors[0] + anchors[1] + anchors[2]) # for all 3 scales44 self.num_anchors = self.anchors.shape[0]45 self.num_anchors_per_scale = self.num_anchors // 346 self.C = C47 self.ignore_iou_thresh = 0.548 49 def __len__(self):50 return len(self.annotations)51 52 def load_mosaic(self, index):53 # YOLOv5 4-mosaic loader. Loads 1 image + 3 random images into a 4-image mosaic54 labels4 = []55 s = self.image_size56 yc, xc = (int(random.uniform(x, 2 * s - x)) for x in self.mosaic_border) # mosaic center x, y57 indices = [index] + random.choices(range(len(self)), k=3) # 3 additional image indices58 random.shuffle(indices)59 for i, index in enumerate(indices):60 # Load image61 label_path = os.path.join(self.label_dir, self.annotations.iloc[index, 1])62 bboxes = np.roll(np.loadtxt(fname=label_path, delimiter=" ", ndmin=2), 4, axis=1).tolist()63 img_path = os.path.join(self.img_dir, self.annotations.iloc[index, 0])64 img = np.array(Image.open(img_path).convert("RGB"))65 66 67 h, w = img.shape[0], img.shape[1]68 labels = np.array(bboxes)69 70 # place img in img471 if i == 0: # top left72 img4 = np.full((s * 2, s * 2, img.shape[2]), 114, dtype=np.uint8) # base image with 4 tiles73 x1a, y1a, x2a, y2a = max(xc - w, 0), max(yc - h, 0), xc, yc # xmin, ymin, xmax, ymax (large image)74 x1b, y1b, x2b, y2b = w - (x2a - x1a), h - (y2a - y1a), w, h # xmin, ymin, xmax, ymax (small image)75 elif i == 1: # top right76 x1a, y1a, x2a, y2a = xc, max(yc - h, 0), min(xc + w, s * 2), yc77 x1b, y1b, x2b, y2b = 0, h - (y2a - y1a), min(w, x2a - x1a), h78 elif i == 2: # bottom left79 x1a, y1a, x2a, y2a = max(xc - w, 0), yc, xc, min(s * 2, yc + h)80 x1b, y1b, x2b, y2b = w - (x2a - x1a), 0, w, min(y2a - y1a, h)81 elif i == 3: # bottom right82 x1a, y1a, x2a, y2a = xc, yc, min(xc + w, s * 2), min(s * 2, yc + h)83 x1b, y1b, x2b, y2b = 0, 0, min(w, x2a - x1a), min(y2a - y1a, h)84 85 img4[y1a:y2a, x1a:x2a] = img[y1b:y2b, x1b:x2b] # img4[ymin:ymax, xmin:xmax]86 padw = x1a - x1b87 padh = y1a - y1b88 89 # Labels90 if labels.size:91 labels[:, :-1] = xywhn2xyxy(labels[:, :-1], w, h, padw, padh) # normalized xywh to pixel xyxy format92 labels4.append(labels)93 94 # Concat/clip labels95 labels4 = np.concatenate(labels4, 0)96 for x in (labels4[:, :-1],):97 np.clip(x, 0, 2 * s, out=x) # clip when using random_perspective()98 # img4, labels4 = replicate(img4, labels4) # replicate99 labels4[:, :-1] = xyxy2xywhn(labels4[:, :-1], 2 * s, 2 * s)100 labels4[:, :-1] = np.clip(labels4[:, :-1], 0, 1)101 labels4 = labels4[labels4[:, 2] > 0]102 labels4 = labels4[labels4[:, 3] > 0]103 return img4, labels4 104 105 def __getitem__(self, index):106 107 image, bboxes = self.load_mosaic(index)108 109 if self.transform:110 augmentations = self.transform(image=image, bboxes=bboxes)111 image = augmentations["image"]112 bboxes = augmentations["bboxes"]113 114 # Below assumes 3 scale predictions (as paper) and same num of anchors per scale115 targets = [torch.zeros((self.num_anchors // 3, S, S, 6)) for S in self.S]116 for box in bboxes:117 iou_anchors = iou(torch.tensor(box[2:4]), self.anchors)118 anchor_indices = iou_anchors.argsort(descending=True, dim=0)119 x, y, width, height, class_label = box120 has_anchor = [False] * 3 # each scale should have one anchor121 for anchor_idx in anchor_indices:122 scale_idx = anchor_idx // self.num_anchors_per_scale123 anchor_on_scale = anchor_idx % self.num_anchors_per_scale124 S = self.S[scale_idx]125 i, j = int(S * y), int(S * x) # which cell126 anchor_taken = targets[scale_idx][anchor_on_scale, i, j, 0]127 if not anchor_taken and not has_anchor[scale_idx]:128 targets[scale_idx][anchor_on_scale, i, j, 0] = 1129 x_cell, y_cell = S * x - j, S * y - i # both between [0,1]130 width_cell, height_cell = (131 width * S,132 height * S,133 ) # can be greater than 1 since it's relative to cell134 box_coordinates = torch.tensor(135 [x_cell, y_cell, width_cell, height_cell]136 )137 targets[scale_idx][anchor_on_scale, i, j, 1:5] = box_coordinates138 targets[scale_idx][anchor_on_scale, i, j, 5] = int(class_label)139 has_anchor[scale_idx] = True140 141 elif not anchor_taken and iou_anchors[anchor_idx] > self.ignore_iou_thresh:142 targets[scale_idx][anchor_on_scale, i, j, 0] = -1 # ignore prediction143 144 return image, tuple(targets)145 146 147def test():148 anchors = config.ANCHORS149 150 transform = config.test_transforms151 152 dataset = YOLODataset(153 "COCO/train.csv",154 "COCO/images/images/",155 "COCO/labels/labels_new/",156 S=[13, 26, 52],157 anchors=anchors,158 transform=transform,159 )160 S = [13, 26, 52]161 scaled_anchors = torch.tensor(anchors) / (162 1 / torch.tensor(S).unsqueeze(1).unsqueeze(1).repeat(1, 3, 2)163 )164 loader = DataLoader(dataset=dataset, batch_size=1, shuffle=True)165 for x, y in loader:166 boxes = []167 168 for i in range(y[0].shape[1]):169 anchor = scaled_anchors[i]170 print(anchor.shape)171 print(y[i].shape)172 boxes += cells_to_bboxes(173 y[i], is_preds=False, S=y[i].shape[2], anchors=anchor174 )[0]175 boxes = nms(boxes, iou_threshold=1, threshold=0.7, box_format="midpoint")176 print(boxes)177 plot_image(x[0].permute(1, 2, 0).to("cpu"), boxes)178 179 180if __name__ == "__main__":181 test()