Spaces:
Sleeping
Sleeping
| 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) | |
| hf_token = os.environ.get("HF_TOKEN") | |
| if hf_token: | |
| login(token=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" # 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 | |
| model = tf.keras.models.load_model(model_path) | |
| # Streamlit UI | |
| tab1, tab2 = st.tabs(["Christmas Tree Classifier", "Sample Image Links"]) | |
| # Tab 1: Christmas Tree Classifier | |
| with tab1: | |
| 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) | |
| st.image(image, caption="Uploaded Image.", use_container_width=True) # Updated to use_container_width | |
| 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}") | |
| # Tab 2: Sample Image Links | |
| with tab2: | |
| st.title("Sample Image Links") | |
| # First header and paragraph with link | |
| st.header("View some of my decorated and undecorated tree samples for the Model here:") | |
| st.write("You can view the sample images for both decorated and undecorated trees in the following link:") | |
| st.write("[View sample images here](https://www.dropbox.com/scl/fo/cuzo12z39cxv6joz7gz2o/ACf5xSjT7nHqMRdgh21GYlc?rlkey=w10usqhkngf2uxwvllgnqb8tf&st=ld22fl4c&dl=0)") | |
| # Second header and paragraph with download link | |
| st.header("Download the tree sample pictures to test them on the model yourself here:") | |
| st.write("You can download the tree sample images by clicking on the link below to test them in the Christmas Tree Classifier:") | |
| st.write("[Download sample images here](https://www.dropbox.com/scl/fo/cuzo12z39cxv6joz7gz2o/ACf5xSjT7nHqMRdgh21GYlc?rlkey=w10usqhkngf2uxwvllgnqb8tf&st=ld22fl4c&dl=1)") |