| 122 | self.p.terminate() |
| 123 | |
| 124 | def run(self): |
| 125 | img_list = [] |
| 126 | is_label_index = True |
| 127 | for line in open(self.img_list_file, 'r'): |
| 128 | item = line.strip('\n').split(self.delimiter) |
| 129 | if len(item) < 2: |
| 130 | is_label_index = False |
| 131 | img_list.append((item[0].strip(), None)) |
| 132 | else: |
| 133 | if not item[1].strip().isdigit(): |
| 134 | # the meta info is not label index |
| 135 | is_label_index = False |
| 136 | img_list.append((item[0].strip(), item[1].strip())) |
| 137 | index = 0 # index for the image |
| 138 | if self.shuffle: |
| 139 | random.shuffle(img_list) |
| 140 | while not self.stop: |
| 141 | if not self.queue.full(): |
| 142 | x, y = [], [] |
| 143 | i = 0 |
| 144 | while i < self.batch_size: |
| 145 | img_path, img_meta = img_list[index] |
| 146 | aug_images = self.image_transform( |
| 147 | os.path.join(self.image_folder, img_path)) |
| 148 | assert i + len(aug_images) <= self.batch_size, \ |
| 149 | 'too many images (%d) in a batch (%d)' % \ |
| 150 | (i + len(aug_images), self.batch_size) |
| 151 | for img in aug_images: |
| 152 | ary = np.asarray(img.convert('RGB'), dtype=np.float32) |
| 153 | x.append(ary.transpose(2, 0, 1)) |
| 154 | if is_label_index: |
| 155 | y.append(int(img_meta)) |
| 156 | else: |
| 157 | y.append(img_meta) |
| 158 | i += 1 |
| 159 | index += 1 |
| 160 | if index == self.num_samples: |
| 161 | index = 0 # reset to the first image |
| 162 | if self.shuffle: |
| 163 | random.shuffle(img_list) |
| 164 | # enqueue one mini-batch |
| 165 | if is_label_index: |
| 166 | self.queue.put((np.asarray(x), np.asarray(y, |
| 167 | dtype=np.int32))) |
| 168 | else: |
| 169 | self.queue.put((np.asarray(x), y)) |
| 170 | else: |
| 171 | time.sleep(0.1) |
| 172 | return |
| 173 | |
| 174 | |
| 175 | if __name__ == '__main__': |