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"  # 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
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)
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) ...