Hudda's picture
Update app.py
303526c verified
raw
history blame
2.29 kB
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}")