Update app.py

app.py

@@ -1,11 +1,24 @@
 import streamlit as st
-from transformers import TFAutoModelForImageClassification, AutoFeatureExtractor
+import tensorflow as tf
 from PIL import Image
 import numpy as np
+import tempfile
+import zipfile
+import os
 
-# Load the model and feature extractor
-model = TFAutoModelForImageClassification.from_pretrained('/content/drive/MyDrive/my_christmas_tree_model')
-feature_extractor = AutoFeatureExtractor.from_pretrained(model.config._name_or_path)
+# Function to load model from a zip file
+def load_model_from_zip(zip_file_path):
+    with tempfile.TemporaryDirectory() as temp_dir:
+        with zipfile.ZipFile(zip_file_path, 'r') as zip_ref:
+            zip_ref.extractall(temp_dir)
+        
+        # Load the model (assuming your model is 'your_trained_model.keras' inside the zip)
+        model_path = os.path.join(temp_dir, 'your_trained_model.keras')
+        model = tf.keras.models.load_model(model_path)  
+    return model
+
+# Load the model from the zip file
+model = load_model_from_zip('/content/drive/MyDrive/my_christmas_tree_model.zip') # Replace 'my_christmas_tree_model.zip' with the actual name of the zip file
 
 # Streamlit UI
 st.title("Christmas Tree Classifier")
@@ -19,15 +32,17 @@ if uploaded_file is not None:
     st.image(image, caption="Uploaded Image.", use_column_width=True)
 
     # Preprocess the image
-    inputs = feature_extractor(images=image, return_tensors="tf")
+    image = image.resize((150, 150))  # Resize to match your model's input shape
+    image_array = np.array(image) / 255.0  # Normalize
+    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]
-    
-    # Map class index to label
-    class_names = model.config.id2label  # Get class names from model config
-    predicted_class = class_names[predicted_class_idx]
+    prediction = model.predict(image_array)
+
+    # Interpret prediction (assuming binary classification)
+    class_names = ['Undecorated', 'Decorated'] #Update your class names here
+    predicted_class_index = 1 if prediction[0][0] >= 0.5 else 0  # Adjust threshold if needed
+    predicted_class = class_names[predicted_class_index]
 
     # Display the prediction
     st.write(f"Prediction: **{predicted_class}**")