Update app.py

app.py

@@ -12,16 +12,37 @@ if hf_token:
 
 # Download and load the model from the Hugging Face Hub
 repo_id = os.environ.get("MODEL_ID", "willco-afk/tree-test-x")  # Get repo ID from secret or default
-filename = "your_trained_model.keras"  # Name of your .keras file
-cache_dir = "./models"  # Local directory to cache the model (create if it doesn't exist)
-os.makedirs(cache_dir, exist_ok=True)  # Create the directory if it doesn't exist
+filename = "your_trained_model.keras"  # Updated filename
+cache_dir = "./models"  # Local directory to cache the model
+os.makedirs(cache_dir, exist_ok=True)  
 model_path = hf_hub_download(repo_id=repo_id, filename=filename, cache_dir=cache_dir)
 
-# Load the model (this line should replace the Google Drive loading line)
+# Load the model
 model = tf.keras.models.load_model(model_path) 
 
 # Streamlit UI
 st.title("Christmas Tree Classifier")
 st.write("Upload an image of a Christmas tree to classify it:")
 
-# ... (rest of your Streamlit code remains the same) ...
+uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])
+
+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
+    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
+    prediction = model.predict(image_array)
+    
+    # Get predicted class
+    predicted_class = "Decorated" if prediction[0][0] >= 0.5 else "Undecorated"
+
+    # Display the prediction
+    st.write(f"Prediction: {predicted_class}")