import numpy as np
import streamlit as st
import tensorflow as tf
from PIL import Image
from huggingface_hub import hf_hub_download
import os


class_names = [
 'Tomato_Bacterial_spot',
 'Tomato_Early_blight',
 'Tomato_Late_blight',
 'Tomato_Leaf_Mold',
 'Tomato_Septoria_leaf_spot',
 'Tomato_Spider_mites_Two_spotted_spider_mite',
 'Tomato__Target_Spot',
 'Tomato__Tomato_YellowLeaf__Curl_Virus',
 'Tomato__Tomato_mosaic_virus',
 'Tomato_healthy'
]

@st.cache_resource
def load_model_from_hub():
    model_path = hf_hub_download(
        repo_id="Heizsenberg/leaf_classification",
        filename="leaf_detection_model.keras",
        repo_type="space"   # VERY IMPORTANT
    )
    model = tf.keras.models.load_model(model_path)
    return model



def predict_image(model, img):
    # Preprocess
    img = img.convert("RGB")  # VERY IMPORTANT
    img = img.resize((128, 128))
    img_array = np.array(img) / 255.0
    img_array = np.expand_dims(img_array, axis=0)

    # 4. Predict
    predictions = model.predict(img_array)
    predicted_index = np.argmax(predictions[0])
    
    predicted_label = class_names[predicted_index]
    
    confidence = np.max(predictions[0])

    return img, predicted_label, confidence

def run():
    classifier_model = load_model_from_hub()
    
    st.write("upload a tomato leaf image to be predicted")
    
    uploaded_file = st.file_uploader(
        "Choose an tomato leaf image to be uploaded",
        type=["JPG", "jpg", "jpeg"] # Specify accepted file types
    )
    
    if uploaded_file is not None:
        st.success("File uploaded successfully!")
        st.write("Filename:", uploaded_file.name)
        st.write(uploaded_file)
        
        st.write("Image")
        # To read image file buffer as a PIL Image:
        img = Image.open(uploaded_file)

        image, predicted_class, confidence = predict_image(
            classifier_model, img
        )

        confidence = confidence * 100
        
        # show result
        st.image(img, caption="Uploaded Image", use_container_width=True)
        st.success(f"Tomato Leaf Prediction: {predicted_class}")
        st.write(f"with confidence level of: {confidence}")

    
if __name__ == '__main__':
    run()