Spaces:
Sleeping
Sleeping
import streamlit as st | |
import tensorflow as tf | |
from PIL import Image | |
import os | |
from huggingface_hub import hf_hub_url, set_access_token # Import Hugging Face utilities | |
# Title of the Streamlit app | |
st.title("Yellow Rust Severity Prediction") | |
# Load Hugging Face API token from environment | |
huggingface_api_token = os.getenv("HUGGINGFACE_TOKEN") | |
if huggingface_api_token is None: | |
st.error("YellowRust API token not found in environment. Please set it.") | |
st.stop() | |
# Set Hugging Face token for authentication | |
set_access_token(huggingface_api_token) | |
# Model repository details | |
model_repo_id = "shaheer-data/Yellow-Rust-Prediction" | |
model_file_path = "final_meta_model.keras" | |
# Construct the model URL | |
st.write("Loading model from Hugging Face repo:", model_repo_id) | |
model_url = hf_hub_url(model_repo_id, model_file_path) | |
loaded_model = tf.keras.models.load_model(model_url) # Load model using tf.keras directly | |
# Function to make predictions | |
def predict_image(image): | |
image = image.resize((224, 224)) # Resize to match model input dimensions | |
image_array = tf.keras.preprocessing.image.img_to_array(image) | |
image_array = tf.expand_dims(image_array, axis=0) # Expand dimensions for batch prediction | |
predictions = loaded_model.predict(image_array) | |
return predictions | |
# Class labels for Yellow Rust severity levels | |
CLASS_LABELS = [ | |
"Healthy", | |
"Mild Severity", | |
"Moderate Severity", | |
"Severe Severity", | |
"Very Severe", | |
"Extreme Severity" | |
] | |
# Image upload widget | |
uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"]) | |
if uploaded_file is not None: | |
image = Image.open(uploaded_file) | |
st.image(image, caption="Uploaded Image", use_column_width=True) | |
# Display progress bar | |
with st.spinner("Making predictions..."): | |
predictions = predict_image(image) | |
predicted_class = predictions.argmax(axis=-1) | |
st.write(f"Predicted Severity Level: {CLASS_LABELS[predicted_class[0]]} with confidence {predictions[0][predicted_class[0]]:.2f}") | |
else: | |
st.write("Please upload an image file to make predictions.") | |