import streamlit as st
from PIL import Image
import os
import torch
from transformers import pipeline  # Assuming you use a Hugging Face pipeline

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

# Ensure the folder exists
if not os.path.exists(FIRE_SHOT_FOLDER):
    st.error("FireShot folder does not exist!")
else:
    st.success(f"FireShot folder found at {FIRE_SHOT_FOLDER}")

# Display a list of image files in the folder
image_files = [f for f in os.listdir(FIRE_SHOT_FOLDER) if f.endswith(('jpg', 'jpeg', 'png'))]

# Make sure there are images in the folder
if not image_files:
    st.info("No images found in FireShot folder.")
else:
    st.write("## Select an image to perform object detection:")

    # Display images as clickable thumbnails
    selected_image = None
    cols = st.columns(4)  # Adjust the number of columns based on your layout preference

    for i, image_file in enumerate(image_files):
        img_path = os.path.join(FIRE_SHOT_FOLDER, image_file)
        img = Image.open(img_path)

        # Display thumbnail of the image
        with cols[i % 4]:  # Display images in 4 columns
            if st.button(image_file):
                selected_image = img_path  # Set the selected image
    
    # Once an image is selected, run the object detection model on it
    if selected_image:
        st.write(f"### Running object detection on {os.path.basename(selected_image)}")
        
        img = Image.open(selected_image)
        st.image(img, caption="Selected Image", use_column_width=True)
        
        # Load your object detection model from Hugging Face (replace with your model)
        model = pipeline('object-detection', model='your-huggingface-model')  # Load your Hugging Face model
        
        # Convert the image to the format needed by the model (PIL Image)
        results = model(selected_image)  # Pass image file to model
        
        # Display results (adjust this depending on the format of model's output)
        st.write(f"### Detected Objects: {results}")
        
        # Assuming you're counting trees based on detection
        tree_count = sum(1 for obj in results if obj['label'] == 'tree')
        st.success(f"### Number of Trees: {tree_count}")