Spaces:
Sleeping
Sleeping
File size: 3,949 Bytes
f29fe4a |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 |
import streamlit as st
from PIL import Image
import cv2
from ultralytics import YOLO
with st.sidebar:
st.title("Control panel")
file = st.file_uploader("Choose an image or a video", type=["png", "jpg", "jpeg", "mp4"])
radio_button1 = st.radio("Model", ["model_train_17", "model_train_15"])
radio_button2=st.radio("Visualize",["No","Yes"])
st.header("Palm Tree Detection")
st.write(
'<p style="font-family: Arial, sans-serif; font-size: px; color: black; font-style: italic;">Counting the number of palm and coconut trees</p>',
unsafe_allow_html=True
)
status_placeholder = st.empty()
if radio_button1 == "model_train_17":
model = YOLO(r'C:\Users\Tectoro\Desktop\Palm tree detection\train_17_best.pt')
elif radio_button1 == "model_train_15":
model = YOLO(r'C:\Users\Tectoro\Desktop\Palm tree detection\train_15_best.pt')
def count_objects(results, class_names):
"""Count objects detected for each class."""
class_counts = {name: 0 for name in class_names.values()}
for box in results[0].boxes:
cls_idx = int(box.cls[0])
class_name = class_names.get(cls_idx, None)
if class_name:
class_counts[class_name] += 1
else:
st.warning(f"Unknown class index detected: {cls_idx}")
return class_counts
def run_inference(file):
file_type = file.type.split('/')[0]
if file_type == 'image':
image = Image.open(file)
st.image(image, caption="Uploaded Image", use_container_width=True)
status_placeholder.write("Processing...Please wait....")
results = model.predict(source=image, save=False)
class_names = model.names
counts = count_objects(results, class_names)
st.write("Detected objects:")
for obj, count in counts.items():
st.write(f"{obj}: {count}")
status_placeholder.empty()
if(radio_button2=="Yes"):
status_placeholder.write("Processing...")
st.image(results[0].plot(), caption="Detected Objects", use_container_width=True)
status_placeholder.empty()
# elif file_type == 'video':
# temp_file = f"temp_{file.name}"
# compressed_file = f"compressed_{file.name}"
# # Save the uploaded video to a temporary file
# with open(temp_file, "wb") as f:
# f.write(file.getbuffer())
# # Compress the video
# st.write("Compressing video...")
# clip = VideoFileClip(temp_file)
# clip.write_videofile(compressed_file, codec="libx264", audio_codec="aac")
# clip.close()
# st.write("Compression complete. Processing video...")
# # Process the compressed video
# cap = cv2.VideoCapture(compressed_file)
# stframe = st.empty()
# total_counts = {name: 0 for name in model.names}
# while cap.isOpened():
# ret, frame = cap.read()
# if not ret:
# break
# # Perform inference on each video frame
# results = model.predict(source=frame, save=False)
# # Count the objects in the frame
# frame_counts = {model.names[int(box.cls[0])]: 0 for box in results[0].boxes}
# for box in results[0].boxes:
# class_name = model.names[int(box.cls[0])]
# frame_counts[class_name] += 1
# for obj, count in frame_counts.items():
# total_counts[obj] += count
# # Display the processed video frame
# stframe.image(results[0].plot(), channels="BGR", use_container_width=True)
# cap.release()
# st.write("Video processing complete.")
# # Display total counts
# st.write("Total detected objects in the video:")
# for obj, count in total_counts.items():
# st.write(f"{obj}: {count}")
if file is not None:
run_inference(file)
|