How can I stop this TFJS model from overfitting? - tensorflow

Question

I am training a simple image classification model with TFJS but the model keeps overfitting. I've experimented with various models/hyperparameters to no avail. I'm beginning to thing the problem is with the dataset itself, not the model as all models/params I've tried work perfectly on different datasets.
About the data itself, I have a small dataset of 600 flower images and a (3 actually) csv file containing the image paths and their labels that looks like this
path
label
05_067.png
9
05_068.png
2
...
...
The code below loads the CSV file, access the images through their paths and converts them to tensors.
/** load and normalize data */
const loadData = function (dataUrl: string, batches = batchSize) {
/** transform input array (xs) to 3D tensor, binarize output label (ys) */
const transform = ({ xs, ys }: any) => {
/** array of numbes (0 - labels.lengh) for use in `one-hot-ing` label values */
const zeros = [...Array(labels.length).keys()];
let buffer = fs.readFileSync(path.join(directory, `${xs[0]}`)),
imageTensor = tf.node.decodeImage(buffer, 3)
.resizeNearestNeighbor([128, 128])
.toFloat()
.div(tf.scalar(225.0)) //dividing pixel values by 225, converting the tensor to a Float32 datatype
// .expandDims()
return {
xs: imageTensor,
ys: tf.tensor1d(zeros.map(i => (i === ys ? 1 : 0)))
};
};
// load, normalize, transform, batch
return tf.data.csv(dataUrl, { columnConfigs: { label: { isLabel: true } } })
.map(({ xs, ys }: any) => ({ xs: Object.values(xs), ys: ys.label }))
.map(transform)
.batch(batches)
};
And this is the model that performs the highest on this dataset (0.003 training l0ss, 0.999 training accuracy, 1.9 validation loss and 0.5 validation accuracy).
/** Define the model architecture */
const model = tf.sequential();
model.add(tf.layers.conv2d({
inputShape: [128, 128, 3],
filters: 64,
kernelSize: 3,
activation: 'relu',
}));
model.add(tf.layers.maxPooling2d({ poolSize: 2, strides: 2 }));
model.add(tf.layers.dropout({ rate: 0.25 }));
model.add(tf.layers.conv2d({
filters: 128,
kernelSize: 3,
activation: 'relu',
}));
model.add(tf.layers.maxPooling2d({ poolSize: 2, strides: 2 }));
model.add(tf.layers.dropout({ rate: 0.25 }));
model.add(tf.layers.flatten());
model.add(tf.layers.dense({ units: 64, activation: 'relu' }));
model.add(tf.layers.dense({ units: labels.length, activation: 'softmax' }));
model.compile({
optimizer: tf.train.adam(0.001),
loss: 'categoricalCrossentropy',
metrics: ['accuracy'],
});
export default model;
This model goes up to 0.96 validation accuracy with the MNIST dataset but then that dataset contains pixel values from the start, not image paths.
My question now is 'What exactly am i doing wrong here?'. Am i loading the data wrongly or am i just using the wrong models?
I would like to add that i have the same dataset in another project but structured as individual folders for each label, and that trains and validates perfectly with the model above.