Update app.py

app.py

@@ -1,36 +1,72 @@
-import os import streamlit as st import tensorflow as tf from PIL import Image import numpy as np from huggingface_hub import login, hf_hub_download
-Authenticate with Hugging Face token (if available)
+import os
+import streamlit as st
+import tensorflow as tf
+from PIL import Image
+import numpy as np
+from huggingface_hub import login, hf_hub_download
 
-hf_token = os.environ.get("HF_TOKEN") if hf_token: login(token=hf_token)
-Download and load the model from the Hugging Face Hub
+# Authenticate with Hugging Face token (if available)
+hf_token = os.environ.get("HF_TOKEN")
+if hf_token:
+    login(token=hf_token)
 
-repo_id = os.environ.get("MODEL_ID", "willco-afk/tree-test-x") # Get repo ID from secret or default filename = "your_trained_model.keras" # Updated filename cache_dir = "./models" # Local directory to cache the model os.makedirs(cache_dir, exist_ok=True)
+# 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_resnet50.keras.zip"  # Updated filename
+cache_dir = "./models"  # Local directory to cache the model
+os.makedirs(cache_dir, exist_ok=True)
+
+# Download the model file from Hugging Face
 model_path = hf_hub_download(repo_id=repo_id, filename=filename, cache_dir=cache_dir)
-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:")
-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)
-# Updated Line:
-st.image(image, caption="Uploaded Image.", use_container_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}")
+
+# Extract and load the model
+model_unzipped_path = os.path.join(cache_dir, "your_trained_model_resnet50.keras")  # Path where we will extract the model
+if not os.path.exists(model_unzipped_path):
+    with zipfile.ZipFile(model_path, 'r') as zip_ref:
+        zip_ref.extractall(cache_dir)
+    print(f"Model unzipped to {model_unzipped_path}")
+
+# Load the model
+model = tf.keras.models.load_model(model_unzipped_path)
+
+# Function for image prediction
+def predict_decoration(image: Image.Image):
+    # Preprocess the image to match the model input size
+    image = image.resize((224, 224))  # Resize to match model's expected input size
+    image_array = np.array(image) / 255.0  # Normalize image to [0, 1]
+    image_array = np.expand_dims(image_array, axis=0)  # Add batch dimension
+
+    # Make prediction
+    prediction = model.predict(image_array)
+    return "Decorated" if prediction[0][0] >= 0.5 else "Undecorated"
+
+# Streamlit UI
+st.title("🎄 Christmas Tree Classifier 🎄")
+st.write("Upload an image of a Christmas tree to classify it:")
+
+# Create tabs
+tab1, tab2 = st.tabs(["Christmas Tree Classifier", "Sample Images"])
+
+# Tab 1: Christmas Tree Classifier
+with tab1:
+    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_container_width=True)
+
+        st.write("Classifying...")
+
+        # Get prediction
+        predicted_class = predict_decoration(image)
+        st.write(f"Prediction: {predicted_class}")
+
+# Tab 2: Sample Images (with text and links only)
+with tab2:
+    st.header("Sample Images for the Model")
+    st.write("View some of my decorated and undecorated tree samples for the Model here:")
+    st.write("[Dropbox Link for Viewing Samples](https://www.dropbox.com/scl/fo/cuzo12z39cxv6joz7gz2o/h?rlkey=w10usqhkngf2uxwvllgnqb8tf&st=ld22fl4c&dl=0)")
+
+    st.write("Download the tree sample images to test them on the model yourself here:")
+    st.write("[Dropbox Link for Downloading Samples](https://www.dropbox.com/scl/fo/cuzo12z39cxv6joz7gz2o/h?rlkey=w10usqhkngf2uxwvllgnqb8tf&st=ld22fl4c&dl=1)")