This commit is contained in:
mszmyd 2024-05-05 19:03:15 +02:00
parent 8e6318b1fe
commit 1cfb74db6a

View File

@ -31,12 +31,12 @@ class Dataset:
self.shuffle_buffer_size = shuffle_buffer_size self.shuffle_buffer_size = shuffle_buffer_size
self.batch_size = batch_size self.batch_size = batch_size
self.dataset = self._load_dataset()\ self.dataset = self.__load_dataset()\
.shuffle(self.shuffle_buffer_size, seed=self.seed)\ .shuffle(self.shuffle_buffer_size, seed=self.seed)\
.repeat(self.repeat)\ .repeat(self.repeat)\
.prefetch(tf.data.experimental.AUTOTUNE) .prefetch(tf.data.experimental.AUTOTUNE)
def _load_dataset(self) -> tf.data.Dataset: def __load_dataset(self) -> tf.data.Dataset:
# check if path has 'test' word in it # check if path has 'test' word in it
dataset = tf.data.Dataset.list_files(str(self.data_dir / '*/*')) dataset = tf.data.Dataset.list_files(str(self.data_dir / '*/*'))
if 'test' in str(self.data_dir).lower(): if 'test' in str(self.data_dir).lower():
@ -44,11 +44,11 @@ class Dataset:
pass pass
else: else:
dataset = dataset.map( dataset = dataset.map(
self._preprocess, num_parallel_calls=tf.data.experimental.AUTOTUNE) self.__preprocess, num_parallel_calls=tf.data.experimental.AUTOTUNE)
return dataset return dataset
def _get_labels(self, image_path): def __get_labels(self, image_path):
path = tf.strings.split(image_path, os.path.sep)[-2] path = tf.strings.split(image_path, os.path.sep)[-2]
plant = tf.strings.split(path, '___')[0] plant = tf.strings.split(path, '___')[0]
disease = tf.strings.split(path, '___')[1] disease = tf.strings.split(path, '___')[1]
@ -58,14 +58,14 @@ class Dataset:
return tf.cast(one_hot_plant, dtype=tf.uint8, name=None), tf.cast(one_hot_disease, dtype=tf.uint8, name=None) return tf.cast(one_hot_plant, dtype=tf.uint8, name=None), tf.cast(one_hot_disease, dtype=tf.uint8, name=None)
def _get_image(self, image_path): def __get_image(self, image_path):
img = tf.io.read_file(image_path) img = tf.io.read_file(image_path)
img = tf.io.decode_jpeg(img, channels=3) img = tf.io.decode_jpeg(img, channels=3)
return tf.cast(img, dtype=tf.float32, name=None) / 255. return tf.cast(img, dtype=tf.float32, name=None) / 255.
def _preprocess(self, image_path): def __preprocess(self, image_path):
labels = self._get_labels(image_path) labels = self.__get_labels(image_path)
image = self._get_image(image_path) image = self.__get_image(image_path)
# returns X, Y1, Y2 # returns X, Y1, Y2
return image, labels[0], labels[1] return image, labels[0], labels[1]