Update app.py

warning removed ~ considering additional UI changes and image presentations

app.py

@@ -1,67 +1,42 @@
 import streamlit as st
 import os
-from PIL import Image
-import tensorflow as tf
 import numpy as np
-import requests
-from io import BytesIO
+import tensorflow as tf
+from PIL import Image
+import io
+import warnings
 
-# Load the trained Keras model (ensure the model file is available in your Hugging Face Space)
-model = tf.keras.models.load_model("your_trained_model.keras")
+# Filter warnings (done before loading model and processing)
+warnings.filterwarnings("ignore", category=DeprecationWarning)
 
-# Function to predict if the tree is decorated or undecorated
-def predict_image(image):
-    image = image.resize((224, 224))  # Adjust size as required by your model
-    image_array = np.array(image) / 255.0  # Normalize the image
-    image_array = np.expand_dims(image_array, axis=0)  # Add batch dimension
-    prediction = model.predict(image_array)
-    return "Decorated" if prediction[0][0] > 0.5 else "Undecorated"
-
-# UI Setup
-st.set_page_config(page_title="Tree Decoration Prediction", layout="wide")
-st.title("Tree Decoration Prediction App")
+# Load your pre-trained Keras model (adjust the path as needed)
+MODEL_PATH = "your_trained_model.keras"  # Model file location in Hugging Face Space
+model = tf.keras.models.load_model(MODEL_PATH)
 
-# Layout with two sections for Undecorated and Decorated Trees
-col1, col2 = st.columns(2)
-with col1:
-    st.header("Undecorated Trees")
-with col2:
-    st.header("Decorated Trees")
+# Define the function to predict decoration
+def predict_decoration(image: Image.Image):
+    # Preprocess the image to match the model input format
+    image = image.resize((224, 224))  # Resize to match model's expected input size
+    image_array = np.array(image) / 255.0  # Normalize the image to [0, 1]
+    image_array = np.expand_dims(image_array, axis=0)  # Add batch dimension
 
-# Folder paths (since images are uploaded directly to the root folder)
-undecorated_folder = "./"  # Root folder where undecorated images are uploaded
-decorated_folder = "./"  # Root folder where decorated images are uploaded
+    # Make prediction
+    prediction = model.predict(image_array)
+    # Assuming the model outputs a binary result (decorated or undecorated)
+    return "Decorated" if prediction[0] > 0.5 else "Undecorated"
 
-# Helper function to display images in a grid format
-def display_images_in_grid(image_folder, is_decorated=False):
-    images = os.listdir(image_folder)
-    # Filter out any non-image files
-    images = [img for img in images if img.endswith(('.png', '.jpg', '.jpeg'))]
-    
-    # Create a 2-row grid with 5 columns
-    num_images = len(images)
-    rows = (num_images // 5) + (1 if num_images % 5 else 0)  # Calculate the number of rows
+# Set up Streamlit interface
+st.title("Tree Decoration Predictor")
 
-    for i in range(rows):
-        cols = st.columns(5)
-        for j in range(5):
-            index = i * 5 + j
-            if index < num_images:
-                image_path = os.path.join(image_folder, images[index])
-                image = Image.open(image_path)
-                
-                # Display image as clickable
-                with cols[j]:
-                    st.image(image, caption=images[index], use_column_width=True)
-                    # Add hover effect and clickable functionality
-                    if st.button(f"Predict {images[index]}", key=f"button_{index}"):
-                        result = predict_image(image)
-                        st.write(f"Prediction: {result}")
+# Upload image through Streamlit
+uploaded_image = st.file_uploader("Upload an image of a tree", type=["jpg", "jpeg", "png"])
 
-# Display undecorated tree images
-with col1:
-    display_images_in_grid(undecorated_folder, is_decorated=False)
+if uploaded_image:
+    # Open and display the uploaded image
+    image = Image.open(uploaded_image)
+    st.image(image, caption="Uploaded Tree Image", use_column_width=True)
 
-# Display decorated tree images
-with col2:
-    display_images_in_grid(decorated_folder, is_decorated=True)
+    # Predict decoration when the button is clicked
+    if st.button("Predict Decoration"):
+        prediction = predict_decoration(image)
+        st.write(f"Prediction: {prediction}")