From 2cd6805063e54fc64f6d10827a8ee3fe23182f27 Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Thu, 21 Mar 2019 02:28:30 +0200 Subject: [PATCH] updates --- utils/datasets.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/utils/datasets.py b/utils/datasets.py index 6ca270d1..31fc166b 100755 --- a/utils/datasets.py +++ b/utils/datasets.py @@ -113,7 +113,12 @@ class LoadImagesAndLabels: # for training return self def __getitem__(self, index): - return self.load_images(index, index + 1) + imgs, labels0, img_paths, img_shapes = self.load_images(index, index + 1) + labels0[:,0] = index % self.batch_size + + labels = torch.zeros(100, 6) + labels[:min(len(labels0), 100)] = labels0 # max 100 labels per image + return imgs.squeeze(0), labels, img_paths, img_shapes def __next__(self): self.count += 1 # batches @@ -122,6 +127,7 @@ class LoadImagesAndLabels: # for training ia = self.count * self.batch_size # start index ib = min(ia + self.batch_size, self.nF) # end index + return self.load_images(ia, ib) def load_images(self, ia, ib):