Skip to content

Instantly share code, notes, and snippets.

@AlessandroMondin
Last active December 26, 2022 21:09
Show Gist options
  • Save AlessandroMondin/cd046c3e3e4849315e92a98a82e2f6ca to your computer and use it in GitHub Desktop.
Save AlessandroMondin/cd046c3e3e4849315e92a98a82e2f6ca to your computer and use it in GitHub Desktop.
def __getitem__(self, idx):
img_name = self.annotations.iloc[idx, 0]
tg_height = self.annotations.iloc[idx, 1] if self.rect_training else 640
tg_width = self.annotations.iloc[idx, 2] if self.rect_training else 640
label_path = os.path.join(self.root_directory, "labels", self.annot_folder, img_name[:-4] + ".txt")
# to avoid an annoying "UserWarning: loadtxt: Empty input file"
with warnings.catch_warnings():
warnings.simplefilter("ignore")
labels = np.loadtxt(fname=label_path, delimiter=" ", ndmin=2)
# removing annotations with negative values
labels = labels[np.all(labels >= 0, axis=1), :]
# to avoid negative values
labels[:, 3:5] = np.floor(labels[:, 3:5] * 1000) / 1000
img = np.array(Image.open(os.path.join(self.root_directory, self.fname, img_name)).convert("RGB"))
if self.bboxes_format == "coco":
labels[:, -1] -= 1 # 0-indexing the classes of coco labels (1-80 --> 0-79)
labels = np.roll(labels, axis=1, shift=1)
# normalized coordinates are scale invariant, hence after resizing the img we don't resize labels
labels[:, 1:] = coco_to_yolo_tensors(labels[:, 1:], w0=img.shape[1], h0=img.shape[0])
img = resize_image(img, (int(tg_width), int(tg_height)))
if self.transform:
# albumentations requires bboxes to be (x,y,w,h,class_idx)
batch_n = idx // self.bs
if batch_n % 2 == 0:
self.transform[1].p = 1
else:
self.transform[1].p = 0
augmentations = self.transform(image=img, bboxes=np.roll(labels, axis=1, shift=4))
img = augmentations["image"]
labels = np.array(augmentations["bboxes"])
# loss fx requires bboxes to be (class_idx,x,y,w,h)
if len(labels):
labels = np.roll(labels, axis=1, shift=1)
if self.ultralytics_loss:
labels = torch.from_numpy(labels)
out_bboxes = torch.zeros((labels.shape[0], 6))
if len(labels):
out_bboxes[..., 1:] = labels
img = img.transpose((2, 0, 1))
img = np.ascontiguousarray(img)
return torch.from_numpy(img), out_bboxes if self.ultralytics_loss else labels
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment