File size: 1,780 Bytes
1ea1dc8 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 | import tensorflow as tf
from keras.applications.mobilenet_v2 import preprocess_input
class HAMDataGenerator:
def __init__(self, train_df):
self.size = [224, 224]
self.batch_size = 16
self.classes = train_df.dx.cat.categories.astype(str).tolist()
self.data_augmentation = tf.keras.Sequential(
[
tf.keras.layers.RandomRotation(20 / 360),
tf.keras.layers.RandomTranslation(0.2, 0.2),
tf.keras.layers.RandomZoom(0.2),
tf.keras.layers.RandomFlip("horizontal"),
],
name="augmentation",
)
def _load_and_preprocess(self, filename, label):
image = tf.io.read_file(filename=filename)
image = tf.io.decode_jpeg(image, channels=3)
image = tf.image.resize(image, size=self.size)
image = preprocess_input(image)
label = tf.one_hot(label, depth=len(self.classes))
return image, label
def _apply_augmentation(self, image, label):
image = self.data_augmentation(image, training=True)
return image, label
def flow_from_dataframe(self, df, directory, x_col, y_col, shuffle=False):
filepaths = [f"{directory}/{img}" for img in df[x_col]]
labels = df[y_col].values
label_to_index = {label: idx for idx, label in enumerate(self.classes)}
label_indices = [label_to_index[label] for label in labels]
ds = tf.data.Dataset.from_tensor_slices((filepaths, label_indices))
if shuffle:
ds = ds.shuffle(buffer_size=len(filepaths))
ds = ds.map(self._load_and_preprocess)
ds = ds.batch(batch_size=self.batch_size)
return ds.map(self._apply_augmentation)
|