import streamlit as st
import os
from PIL import Image
from transformers import pipeline  # Hugging Face pipeline for object detection
import time  # To simulate a delay (for model processing)

# Define the path to the folder where images are stored
FIRE_SHOT_FOLDER = os.path.join(os.getcwd(), 'FireShot')


# Get list of images from the FireShot folder
image_files = [f for f in os.listdir(FIRE_SHOT_FOLDER) if f.endswith(('jpg', 'jpeg', 'png'))]

if image_files:
    # Create a selectbox for the user to select an image
    selected_image = st.selectbox("Select an image from the FireShot folder", image_files)

    # Load and display the selected image
    if selected_image:
        img_path = os.path.join(FIRE_SHOT_FOLDER, selected_image)
        img = Image.open(img_path)
        st.image(img, caption=f"Selected Image: {selected_image}", use_column_width=True)

        # Add a button to start the object detection
        if st.button("Run Tree Detection"):
            # Show loading spinner while processing the model
            with st.spinner("Running YOLO model to detect trees..."):
                # Simulate loading time (you can remove this when using an actual model)
                time.sleep(3)  # Simulate delay (this line can be removed)
                
                # Load YOLO model from Hugging Face
                model = pipeline('object-detection', model="blah-blah-treecounter")

                # Perform object detection on the selected image
                results = model(img_path)

                # Simulate counting trees based on the detected objects
                tree_count = sum(1 for obj in results if obj['label'].lower() == 'tree')

            # Display the results in a button to be clicked
            if st.button("Show Results"):
                st.write(f"### Detected Trees: {tree_count}")
                st.json(results)  # Display full detection results in JSON format if needed

else:
    st.info("No images found in FireShot folder.")