How do I implement train_test_split in this code?

If you are loading the dataset from any directory.. use image_dataset_from_directory from tensorflow

And use subset feature

import tensorflow as tf
from tf.keras.utils import image_dataset_from_directory

path="<put path here>"

training_data=image_dataset_from_directory(
   path,
   image_size=(<put your image size here>),
   batch_size=batch_size,
   validation_split=0.2 ,# as per your need
   subset='training'

)
 validation_data=image_dataset_from_directory(
     path,
     image_size=(<put your image size here>),
     batch_size=batch_size,
     validation_split=0.2 ,# as per your need
     subset='validation'

)

You can add other args like labels,color mode and many others inside the " ( )" .. but it should be same on both place..

If you have any more query , refer to this Documentation

If you got satisfied or your query is clear... upvote... Otherwise let me know in Comments