diff --git a/utils/dataloader.py b/utils/dataloader.py index 100ba37ab4f7fe3988ef6eac6fde9c21b2bb7bc3..7e082deccfda9f7f9332ace68f453a2d670687d2 100644 --- a/utils/dataloader.py +++ b/utils/dataloader.py @@ -104,8 +104,9 @@ class SSDDataset(Dataset): n = self.train_batches temp_index = index % n while True: - img, y = self.get_random_data(lines[temp_index], self.image_size[0:2]) + img, y = self.get_random_data(lines[index], self.image_size[0:2]) if len(y)==0: + index = (index + 1) % n continue boxes = np.array(y[:,:4],dtype=np.float32) boxes[:,0] = boxes[:,0]/self.image_size[1] @@ -114,9 +115,9 @@ class SSDDataset(Dataset): boxes[:,3] = boxes[:,3]/self.image_size[0] boxes = np.maximum(np.minimum(boxes,1),0) if ((boxes[:,3]-boxes[:,1])<=0).any() and ((boxes[:,2]-boxes[:,0])<=0).any(): + index = (index + 1) % n continue y = np.concatenate([boxes,y[:,-1:]],axis=-1) - temp_index = (temp_index + 1) % n break img = np.array(img, dtype=np.float32)