| 93 | |
| 94 | |
| 95 | class LoadSegData: |
| 96 | def __init__(self, opt): |
| 97 | self.opt = opt |
| 98 | |
| 99 | def parse_data(self, image, mask): |
| 100 | # read the image from disk, decode it, convert the data type to |
| 101 | # floating point, and resize it |
| 102 | image = tf.io.read_file(image) |
| 103 | image = tf.image.decode_png(image, channels=3) |
| 104 | image = tf.image.convert_image_dtype(image, dtype=tf.float32) |
| 105 | image = tf.image.resize(image, (self.opt["img_size_h"], self.opt["img_size_w"])) |
| 106 | |
| 107 | mask = tf.io.read_file(mask) |
| 108 | mask = tf.image.decode_png(mask, channels=3) |
| 109 | mask = tf.image.resize(mask, (self.opt["img_size_h"], self.opt["img_size_w"])) |
| 110 | |
| 111 | one_hot_map = [] |
| 112 | for colour in list(color_values.values()): |
| 113 | class_map = tf.reduce_all(tf.equal(mask, colour), axis=-1) |
| 114 | one_hot_map.append(class_map) |
| 115 | one_hot_map = tf.stack(one_hot_map, axis=-1) |
| 116 | one_hot_map = tf.cast(one_hot_map, tf.float32) |
| 117 | |
| 118 | # return the image and the label |
| 119 | return image, one_hot_map |
| 120 | |
| 121 | def init(self, images, masks): |
| 122 | ds = tf.data.Dataset.from_tensor_slices((images, masks)) |
| 123 | if self.opt["use_shuffle"]: |
| 124 | ds = ( |
| 125 | ds.shuffle(len(images)) |
| 126 | .map(self.parse_data, num_parallel_calls=AUTOTUNE) |
| 127 | .batch(self.opt["batch_size"], drop_remainder=True) |
| 128 | .prefetch(AUTOTUNE) |
| 129 | ) |
| 130 | else: |
| 131 | ds = ( |
| 132 | ds.map(self.parse_data, num_parallel_calls=AUTOTUNE) |
| 133 | .batch(self.opt["batch_size"], drop_remainder=True) |
| 134 | .prefetch(AUTOTUNE) |
| 135 | ) |
| 136 | return ds |
| 137 | |
| 138 | |
| 139 | class LoadClassData: |