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

model = load_model('yolo8_model.h5')

def detect_trees(image):
    # Preprocess the image
    img = np.array(image)
    img = img / 255.0

    # Make predictions
    predictions = model.predict(img)

    # Process the predictions
    trees = []
    for prediction in predictions:
        x, y, w, h, confidence = prediction
        if confidence > 0.5:
            trees.append((x, y, w, h))

    return trees

def main():
    st.title('Tree Counter')
    st.write('Upload an image to count the trees')

    col1, col2, col3 = st.columns([1, 1, 1])
    with col1:
        st.button('Detect Trees')
    with col2:
        st.button('Crop Image')
    with col3:
        st.button('Reset')

    uploaded_file = st.file_uploader('Choose an image', type=['jpg', 'png', 'jpeg'])
    if uploaded_file is not None:
        image = Image.open(uploaded_file)
        st.image(image, caption='Uploaded Image')

        width, height = image.size
        if width > 1024 or height > 1024:
            if st.button('Crop Image'):
                # Crop the image
                cropped_image = image.crop((100, 100, 300, 300))
                st.image(cropped_image, caption='Cropped Image')
        else:
            if st.button('Detect Trees'):
                trees = detect_trees(image)
                st.write(f'Number of Trees: {len(trees)}')

    st.write('Google Earth Integration')
    st.map()

if __name__ == '__main__':
    main()