Update app.py

app.py

@@ -1,21 +1,11 @@
 import os
 import streamlit as st
-from transformers import TFAutoModelForImageClassification, AutoFeatureExtractor
+import tensorflow as tf
 from PIL import Image
 import numpy as np
-from huggingface_hub import login
-
-
-# Authenticate with Hugging Face token (if available)
-hf_token = os.environ.get("HF_TOKEN") 
-if hf_token:
-    login(token=hf_token)
-
-
-# Load the model and feature extractor
-model = TFAutoModelForImageClassification.from_pretrained(os.environ.get("MODEL_ID", "willco-afk/tree-test-x")) 
-feature_extractor = AutoFeatureExtractor.from_pretrained(model.config._name_or_path)
 
+# Load your Keras model from Google Drive
+model = tf.keras.models.load_model('/content/drive/MyDrive/your_trained_model.keras')
 
 # Streamlit UI
 st.title("Christmas Tree Classifier")
@@ -27,17 +17,19 @@ if uploaded_file is not None:
     # Display the uploaded image
     image = Image.open(uploaded_file)
     st.image(image, caption="Uploaded Image.", use_column_width=True)
+    st.write("")
+    st.write("Classifying...")
 
     # Preprocess the image
-    inputs = feature_extractor(images=image, return_tensors="tf")
+    image = image.resize((224, 224))  # Resize to match your model's input size
+    image_array = np.array(image) / 255.0  # Normalize pixel values
+    image_array = np.expand_dims(image_array, axis=0)  # Add batch dimension
 
     # Make prediction
-    logits = model(**inputs).logits
-    predicted_class_idx = tf.math.argmax(logits, axis=-1)[0]
+    prediction = model.predict(image_array)
     
-    # Map class index to label
-    class_names = model.config.id2label  # Get class names from model config
-    predicted_class = class_names[predicted_class_idx]
+    # Get predicted class
+    predicted_class = "Decorated" if prediction[0][0] >= 0.5 else "Undecorated"
 
     # Display the prediction
-    st.write(f"Prediction: **{predicted_class}**")
+    st.write(f"Prediction: {predicted_class}")