Spaces:
Sleeping
Sleeping
| import streamlit as st | |
| import torch | |
| import cv2 | |
| import numpy as np | |
| from PIL import Image | |
| import tempfile | |
| import os | |
| from ultralytics import YOLO | |
| # Load the YOLOv8 model from the local path | |
| model = YOLO('best.pt') # Directly reference the best.pt file in the same directory as app.py | |
| def detect_people(image): | |
| results = model(image) | |
| return results | |
| def process_results(results, image): | |
| annotated_frame = image.copy() | |
| person_detected = False | |
| for result in results: | |
| boxes = result.boxes.xyxy.cpu().numpy().astype(int) # Bounding boxes | |
| if len(boxes) > 0: | |
| person_detected = True | |
| for box in boxes: | |
| cv2.rectangle(annotated_frame, (box[0], box[1]), (box[2], box[3]), (0, 255, 0), 2) | |
| if person_detected: | |
| cv2.putText(annotated_frame, 'Canh bao! Co nguoi! Co nguoi', (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 2, cv2.LINE_AA) | |
| return annotated_frame | |
| def main(): | |
| st.title("People Detection Application using YOLOv8") | |
| option = st.sidebar.selectbox("Choose input type", ("Upload Image/Video", "Use Camera")) | |
| if option == "Upload Image/Video": | |
| uploaded_file = st.file_uploader("Upload Image or Video", type=['jpg', 'jpeg', 'png', 'mp4', 'avi', 'mov']) | |
| if uploaded_file is not None: | |
| if uploaded_file.type.startswith('image'): | |
| file_bytes = np.asarray(bytearray(uploaded_file.read()), dtype=np.uint8) | |
| image = cv2.imdecode(file_bytes, 1) | |
| st.image(image, caption='Uploaded Image.', use_column_width=True) | |
| results = detect_people(image) | |
| processed_frame = process_results(results, image) | |
| st.image(processed_frame, caption='Processed Image.', use_column_width=True) | |
| elif uploaded_file.type.startswith('video'): | |
| tfile = tempfile.NamedTemporaryFile(delete=False, suffix='.mp4') | |
| tfile.write(uploaded_file.read()) | |
| tfile.close() | |
| st.write(f'Temporary file path: {tfile.name}') | |
| vidcap = cv2.VideoCapture(tfile.name) | |
| stframe = st.empty() | |
| while vidcap.isOpened(): | |
| success, frame = vidcap.read() | |
| if not success: | |
| break | |
| results = detect_people(frame) | |
| processed_frame = process_results(results, frame) | |
| stframe.image(processed_frame, channels="BGR") | |
| vidcap.release() | |
| os.remove(tfile.name) | |
| elif option == "Use Camera": | |
| stframe = st.empty() | |
| cap = cv2.VideoCapture(0) | |
| while True: | |
| ret, frame = cap.read() | |
| if not ret: | |
| break | |
| results = detect_people(frame) | |
| processed_frame = process_results(results, frame) | |
| stframe.image(processed_frame, channels="BGR") | |
| cap.release() | |
| if __name__ == '__main__': | |
| main() | |