import streamlit as st
import tensorflow as tf
import numpy as np
from PIL import Image
import json

# Load class indices
with open("class_indices.json", "r") as f:
    class_indices = json.load(f)

# Reverse the mapping for predictions
class_names = {v: k for k, v in class_indices.items()}

# Load the TFLite model
interpreter = tf.lite.Interpreter(model_path="model.tflite")
interpreter.allocate_tensors()

# Get input and output details
input_details = interpreter.get_input_details()
output_details = interpreter.get_output_details()

# Define the image preprocessing function
def preprocess_image(image, target_size=(224, 224)):
    image = image.resize(target_size)
    image = np.array(image) / 255.0  # Normalize the image
    image = np.expand_dims(image, axis=0)  # Add batch dimension
    return image.astype(np.float32)

# Define prediction function
def predict(image):
    input_data = preprocess_image(image)
    interpreter.set_tensor(input_details[0]['index'], input_data)
    interpreter.invoke()
    output_data = interpreter.get_tensor(output_details[0]['index'])
    predicted_class = np.argmax(output_data)
    confidence = np.max(output_data)
    return class_names[predicted_class], confidence

# Streamlit UI
st.title("🌾 Crop Disease Prediction")
st.write("Upload an image of a crop leaf, and the app will predict the disease (if any).")

uploaded_file = st.file_uploader("Choose an image file", type=["jpg", "png", "jpeg"])

if uploaded_file is not None:
    # Display the uploaded image
    image = Image.open(uploaded_file)
    st.image(image, caption="Uploaded Image", use_column_width=True)
    
    st.write("Processing...")
    
    # Perform prediction
    predicted_class, confidence = predict(image)
    st.write(f"**Prediction:** {predicted_class}")
    st.write(f"**Confidence:** {confidence:.2f}")