Utility for iterating over an image dataset to get mini-batches. Args: img_list_file(str): name of the file containing image meta data; each line consists of image_path_suffix delimiter meta_info, where meta info could be label ind
| 62 | |
| 63 | |
| 64 | class ImageBatchIter(object): |
| 65 | '''Utility for iterating over an image dataset to get mini-batches. |
| 66 | |
| 67 | Args: |
| 68 | img_list_file(str): name of the file containing image meta data; each |
| 69 | line consists of image_path_suffix delimiter meta_info, |
| 70 | where meta info could be label index or label strings, etc. |
| 71 | meta_info should not contain the delimiter. If the meta_info |
| 72 | of each image is just the label index, then we will parse the |
| 73 | label index into a numpy array with length=batchsize |
| 74 | (for compatibility); otherwise, we return a list of meta_info; |
| 75 | if meta info is available, we return a list of None. |
| 76 | batch_size(int): num of samples in one mini-batch |
| 77 | image_transform: a function for image augmentation; it accepts the full |
| 78 | image path and outputs a list of augmented images. |
| 79 | shuffle(boolean): True for shuffling images in the list |
| 80 | delimiter(char): delimiter between image_path_suffix and label, e.g., |
| 81 | space or comma |
| 82 | image_folder(boolean): prefix of the image path |
| 83 | capacity(int): the max num of mini-batches in the internal queue. |
| 84 | ''' |
| 85 | |
| 86 | def __init__(self, |
| 87 | img_list_file, |
| 88 | batch_size, |
| 89 | image_transform, |
| 90 | shuffle=True, |
| 91 | delimiter=' ', |
| 92 | image_folder=None, |
| 93 | capacity=10): |
| 94 | self.img_list_file = img_list_file |
| 95 | self.queue = Queue(capacity) |
| 96 | self.batch_size = batch_size |
| 97 | self.image_transform = image_transform |
| 98 | self.shuffle = shuffle |
| 99 | self.delimiter = delimiter |
| 100 | self.image_folder = image_folder |
| 101 | self.stop = False |
| 102 | self.p = None |
| 103 | with open(img_list_file, 'r') as fd: |
| 104 | self.num_samples = len(fd.readlines()) |
| 105 | |
| 106 | def start(self): |
| 107 | self.p = Process(target=self.run) |
| 108 | self.p.start() |
| 109 | return |
| 110 | |
| 111 | def __next__(self): |
| 112 | assert self.p is not None, 'call start before next' |
| 113 | while self.queue.empty(): |
| 114 | time.sleep(0.1) |
| 115 | x, y = self.queue.get() # dequeue one mini-batch |
| 116 | return x, y |
| 117 | |
| 118 | def end(self): |
| 119 | if self.p is not None: |
| 120 | self.stop = True |
| 121 | time.sleep(0.1) |