Hudda's picture
Create main.py
77e1798 verified
raw
history blame
1.54 kB
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()