import streamlit as st
from tensorflow.keras.models import load_model
import cv2
import numpy as np
from PIL import Image

# Load the model
model = load_model('best_model.keras')

def preprocess_image(image):
    # Resize the image as required by the model
    img = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
    img = cv2.resize(img, (220, 220))  # Resize to match the model input
    # Normalize the image
    img = img.astype('float32') / 255.0
    # Add a batch dimension
    img = np.expand_dims(img, axis=0)
    return img

def run():
    # Create title
    st.title('Detecting Fire in Forest Images')

    # Create a form for image input
    with st.form('form_forest_fire_detection'):
        # Image upload
        uploaded_image = st.file_uploader('Upload an image', type=['jpg', 'jpeg', 'png'])
        
        # Submit button
        submitted = st.form_submit_button('Detect Fire or No Fire')

    if uploaded_image:
        # Display the uploaded image
        st.image(uploaded_image, caption='Uploaded Image', use_column_width=True)

        if submitted:
            # Preprocess the image
            image = Image.open(uploaded_image)
            preprocessed_image = preprocess_image(image)

            # Predict using the model
            prediction = model.predict(preprocessed_image)
            
            # For example, if prediction > 0.5 classify as 'No Fire', otherwise 'Fire'
            fire_probability = prediction[0][0]
            result = 'No Fire' if fire_probability > 0.5 else 'Fire'

            st.write('## Prediction: ', result)
            st.write('## Raw Prediction Output: ', prediction)


if __name__ == '__main__':
    run()