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

# Function to load the model
@st.cache_resource
def load_model():
    model_path = 'models/my_model.h5'  # Path relative to the script
    if not os.path.isfile(model_path):
        st.error(f"Model file not found: {model_path}")
        return None
    try:
        model = tf.keras.models.load_model(model_path)
        st.success("Model loaded successfully!")
        return model
    except Exception as e:
        st.error(f"Error loading model: {e}")
        return None

# Function to preprocess the image
def preprocess_image(image):
    image = np.array(image.convert('RGB'))
    image = cv2.resize(image, (224, 224))
    image = image / 255.0
    image = np.expand_dims(image, axis=0)
    return image

# Function to predict the class
def predict(image, model):
    processed_image = preprocess_image(image)
    prediction = model.predict(processed_image)
    return prediction

# Main app
def main():
    st.title("Food Item Recognition and Estimation")
    st.write("Upload an image of a food item and the model will recognize the food item and estimate its calories.")

    model = load_model()

    if model is None:
        st.write("Model could not be loaded.")
        return

    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)
        
        st.write("")
        st.write("Classifying...")
        try:
            prediction = predict(image, model)
            predicted_class = np.argmax(prediction, axis=1)[0]
            st.write(f"Predicted class: {predicted_class}")
        except Exception as e:
            st.error(f"Error in prediction: {e}")

if __name__ == "__main__":
    main()