| 83 | self.if_shuffled = False |
| 84 | |
| 85 | def _get_sub_batch(self): |
| 86 | while True: |
| 87 | # get a sample record |
| 88 | this_sample = self.list_sample[self.cur_idx] |
| 89 | if this_sample['height'] > this_sample['width']: |
| 90 | self.batch_record_list[0].append(this_sample) # h > w, go to 1st class |
| 91 | else: |
| 92 | self.batch_record_list[1].append(this_sample) # h <= w, go to 2nd class |
| 93 | |
| 94 | # update current sample pointer |
| 95 | self.cur_idx += 1 |
| 96 | if self.cur_idx >= self.num_sample: |
| 97 | self.cur_idx = 0 |
| 98 | np.random.shuffle(self.list_sample) |
| 99 | |
| 100 | if len(self.batch_record_list[0]) == self.batch_per_gpu: |
| 101 | batch_records = self.batch_record_list[0] |
| 102 | self.batch_record_list[0] = [] |
| 103 | break |
| 104 | elif len(self.batch_record_list[1]) == self.batch_per_gpu: |
| 105 | batch_records = self.batch_record_list[1] |
| 106 | self.batch_record_list[1] = [] |
| 107 | break |
| 108 | return batch_records |
| 109 | |
| 110 | def __getitem__(self, index): |
| 111 | # NOTE: random shuffle for the first time. shuffle in __init__ is useless |