| 137 | |
| 138 | |
| 139 | class LoadClassData: |
| 140 | def __init__(self, opt, num_classes): |
| 141 | self.opt = opt |
| 142 | self.num_classes = num_classes |
| 143 | |
| 144 | def parse_data(self, image, label): |
| 145 | # read the image from disk, decode it, convert the data type to |
| 146 | # floating point, and resize it |
| 147 | image = tf.io.read_file(image) |
| 148 | image = tf.image.decode_png(image, channels=3) |
| 149 | image = tf.image.convert_image_dtype(image, dtype=tf.float32) |
| 150 | image = tf.image.resize(image, (self.opt["img_size_h"], self.opt["img_size_w"])) |
| 151 | |
| 152 | label = label |
| 153 | |
| 154 | # return the image and the label |
| 155 | return image, label |
| 156 | |
| 157 | def init(self, images, labels): |
| 158 | labels = tf.keras.utils.to_categorical(labels, num_classes=self.num_classes) |
| 159 | ds = tf.data.Dataset.from_tensor_slices((images, labels)) |
| 160 | if self.opt["use_shuffle"]: |
| 161 | ds = ( |
| 162 | ds.shuffle(len(images)) |
| 163 | .map(self.parse_data, num_parallel_calls=AUTOTUNE) |
| 164 | .batch(self.opt["batch_size"], drop_remainder=True) |
| 165 | .prefetch(AUTOTUNE) |
| 166 | ) |
| 167 | else: |
| 168 | ds = ( |
| 169 | ds.map(self.parse_data, num_parallel_calls=AUTOTUNE) |
| 170 | .batch(self.opt["batch_size"], drop_remainder=True) |
| 171 | .prefetch(AUTOTUNE) |
| 172 | ) |
| 173 | return ds |
| 174 | |
| 175 | |
| 176 | def get_data(opt): |