Update pages/15_TransferLearning_HF.py

pages/15_TransferLearning_HF.py

@@ -29,14 +29,20 @@ model_name = "google/vit-base-patch16-224-in21k"
 feature_extractor = ViTFeatureExtractor.from_pretrained(model_name)
 base_model = TFAutoModelForImageClassification.from_pretrained(model_name, num_labels=2)  # Cats vs Dogs has 2 classes
 
-# Freeze the convolutional base
+# Freeze the base model
 base_model.trainable = False
 
+# Function to extract features using the feature extractor
+def extract_features(images):
+    # Convert images to the expected format for the feature extractor
+    images = [tf.image.convert_image_dtype(image, tf.float32) for image in images]
+    inputs = feature_extractor(images, return_tensors="tf")
+    return inputs["pixel_values"]
+
 # Add custom layers on top
 inputs = tf.keras.Input(shape=(224, 224, 3))
-features = feature_extractor(inputs, return_tensors="tf")["pixel_values"]
-base_output = base_model(features)[0]  # Extract base model output
-x = tf.keras.layers.Flatten()(base_output)
+features = extract_features([inputs])
+x = base_model.vit(inputs).last_hidden_state[:, 0]
 x = tf.keras.layers.Dense(256, activation='relu')(x)
 x = tf.keras.layers.Dropout(0.5)(x)
 outputs = tf.keras.layers.Dense(1, activation='sigmoid')(x)
@@ -83,4 +89,3 @@ if st.button("Train Model"):
 if st.button("Evaluate Model"):
     test_loss, test_acc = model.evaluate(ds_val, verbose=2)
     st.write(f"Validation accuracy: {test_acc}")
-