Update pages/13_TransferLearning.py

pages/13_TransferLearning.py

@@ -1,13 +1,21 @@
 import streamlit as st
 import tensorflow as tf
 from tensorflow.keras import layers, models, applications
-from tensorflow.keras.preprocessing.image import ImageDataGenerator
+import tensorflow_datasets as tfds
 import matplotlib.pyplot as plt
-import numpy as np
 
-# Set dataset paths
-train_dir = 'data/train'
-validation_dir = 'data/validation'
+# Load the dataset
+dataset_name = "cats_vs_dogs"
+(ds_train, ds_val), ds_info = tfds.load(dataset_name, split=['train[:80%]', 'train[80%:]'], with_info=True, as_supervised=True)
+
+# Preprocess the dataset
+def preprocess_image(image, label):
+    image = tf.image.resize(image, (150, 150))
+    image = image / 255.0
+    return image, label
+
+ds_train = ds_train.map(preprocess_image).batch(32).prefetch(tf.data.AUTOTUNE)
+ds_val = ds_val.map(preprocess_image).batch(32).prefetch(tf.data.AUTOTUNE)
 
 # Streamlit app
 st.title("Transfer Learning with VGG16 for Image Classification")
@@ -16,34 +24,6 @@ st.title("Transfer Learning with VGG16 for Image Classification")
 batch_size = st.slider("Batch Size", 16, 128, 32, 16)
 epochs = st.slider("Epochs", 5, 50, 10, 5)
 
-# Data augmentation and preprocessing
-train_datagen = ImageDataGenerator(
-    rescale=1./255,
-    rotation_range=40,
-    width_shift_range=0.2,
-    height_shift_range=0.2,
-    shear_range=0.2,
-    zoom_range=0.2,
-    horizontal_flip=True,
-    fill_mode='nearest'
-)
-
-validation_datagen = ImageDataGenerator(rescale=1./255)
-
-train_generator = train_datagen.flow_from_directory(
-    train_dir,
-    target_size=(150, 150),
-    batch_size=batch_size,
-    class_mode='binary'
-)
-
-validation_generator = validation_datagen.flow_from_directory(
-    validation_dir,
-    target_size=(150, 150),
-    batch_size=batch_size,
-    class_mode='binary'
-)
-
 # Load the pre-trained VGG16 model
 base_model = applications.VGG16(weights='imagenet', include_top=False, input_shape=(150, 150, 3))
 
@@ -70,11 +50,9 @@ model.compile(optimizer='adam',
 if st.button("Train Model"):
     with st.spinner("Training the model..."):
         history = model.fit(
-            train_generator,
-            steps_per_epoch=train_generator.samples // train_generator.batch_size,
+            ds_train,
             epochs=epochs,
-            validation_data=validation_generator,
-            validation_steps=validation_generator.samples // validation_generator.batch_size
+            validation_data=ds_val
         )
         
         st.success("Model training completed!")
@@ -100,5 +78,5 @@ if st.button("Train Model"):
 
 # Evaluate the model
 if st.button("Evaluate Model"):
-    test_loss, test_acc = model.evaluate(validation_generator, verbose=2)
+    test_loss, test_acc = model.evaluate(ds_val, verbose=2)
     st.write(f"Validation accuracy: {test_acc}")