| from ultralytics import YOLO
|
| import streamlit as st
|
| import cv2
|
| import yt_dlp
|
| import settings
|
|
|
|
|
| def load_model(model_path):
|
| """
|
| Loads a YOLO object detection model from the specified model_path.
|
|
|
| Parameters:
|
| model_path (str): The path to the YOLO model file.
|
|
|
| Returns:
|
| A YOLO object detection model.
|
| """
|
| model = YOLO(model_path)
|
| return model
|
|
|
|
|
| def display_tracker_options():
|
| display_tracker = st.radio("Display Tracker", ('Yes', 'No'))
|
| is_display_tracker = True if display_tracker == 'Yes' else False
|
| if is_display_tracker:
|
| tracker_type = st.radio("Tracker", ("bytetrack.yaml", "botsort.yaml"))
|
| return is_display_tracker, tracker_type
|
| return is_display_tracker, None
|
|
|
|
|
| def _display_detected_frames(conf, model, st_frame, image, is_display_tracking=None, tracker=None):
|
| """
|
| Display the detected objects on a video frame using the YOLOv8 model.
|
|
|
| Args:
|
| - conf (float): Confidence threshold for object detection.
|
| - model (YoloV8): A YOLOv8 object detection model.
|
| - st_frame (Streamlit object): A Streamlit object to display the detected video.
|
| - image (numpy array): A numpy array representing the video frame.
|
| - is_display_tracking (bool): A flag indicating whether to display object tracking (default=None).
|
|
|
| Returns:
|
| None
|
| """
|
|
|
|
|
| image = cv2.resize(image, (720, int(720*(9/16))))
|
|
|
|
|
| if is_display_tracking:
|
| res = model.track(image, conf=conf, persist=True, tracker=tracker)
|
| else:
|
|
|
| res = model.predict(image, conf=conf)
|
|
|
|
|
| res_plotted = res[0].plot()
|
| st_frame.image(res_plotted,
|
| caption='Detected Video',
|
| channels="BGR",
|
| use_column_width=True
|
| )
|
|
|
|
|
| def get_youtube_stream_url(youtube_url):
|
| ydl_opts = {
|
| 'format': 'best[ext=mp4]',
|
| 'no_warnings': True,
|
| 'quiet': True
|
| }
|
| with yt_dlp.YoutubeDL(ydl_opts) as ydl:
|
| info = ydl.extract_info(youtube_url, download=False)
|
| return info['url']
|
|
|
|
|
| def play_youtube_video(conf, model):
|
| source_youtube = st.sidebar.text_input("YouTube Video url")
|
| is_display_tracker, tracker = display_tracker_options()
|
|
|
| if st.sidebar.button('Detect Objects'):
|
| if not source_youtube:
|
| st.sidebar.error("Please enter a YouTube URL")
|
| return
|
|
|
| try:
|
| st.sidebar.info("Extracting video stream URL...")
|
| stream_url = get_youtube_stream_url(source_youtube)
|
|
|
| st.sidebar.info("Opening video stream...")
|
| vid_cap = cv2.VideoCapture(stream_url)
|
|
|
| if not vid_cap.isOpened():
|
| st.sidebar.error(
|
| "Failed to open video stream. Please try a different video.")
|
| return
|
|
|
| st.sidebar.success("Video stream opened successfully!")
|
| st_frame = st.empty()
|
| while vid_cap.isOpened():
|
| success, image = vid_cap.read()
|
| if success:
|
| _display_detected_frames(
|
| conf,
|
| model,
|
| st_frame,
|
| image,
|
| is_display_tracker,
|
| tracker
|
| )
|
| else:
|
| break
|
|
|
| vid_cap.release()
|
|
|
| except Exception as e:
|
| st.sidebar.error(f"An error occurred: {str(e)}")
|
|
|
|
|
| def play_rtsp_stream(conf, model):
|
| """
|
| Plays an rtsp stream. Detects Objects in real-time using the YOLOv8 object detection model.
|
|
|
| Parameters:
|
| conf: Confidence of YOLOv8 model.
|
| model: An instance of the `YOLOv8` class containing the YOLOv8 model.
|
|
|
| Returns:
|
| None
|
|
|
| Raises:
|
| None
|
| """
|
| source_rtsp = st.sidebar.text_input("rtsp stream url:")
|
| st.sidebar.caption(
|
| 'Example URL: rtsp://admin:12345@192.168.1.210:554/Streaming/Channels/101')
|
| is_display_tracker, tracker = display_tracker_options()
|
| if st.sidebar.button('Detect Objects'):
|
| try:
|
| vid_cap = cv2.VideoCapture(source_rtsp)
|
| st_frame = st.empty()
|
| while (vid_cap.isOpened()):
|
| success, image = vid_cap.read()
|
| if success:
|
| _display_detected_frames(conf,
|
| model,
|
| st_frame,
|
| image,
|
| is_display_tracker,
|
| tracker
|
| )
|
| else:
|
| vid_cap.release()
|
| break
|
| except Exception as e:
|
| vid_cap.release()
|
| st.sidebar.error("Error loading RTSP stream: " + str(e))
|
|
|
|
|
| def play_webcam(conf, model):
|
| """
|
| Plays a webcam stream. Detects Objects in real-time using the YOLOv8 object detection model.
|
|
|
| Parameters:
|
| conf: Confidence of YOLOv8 model.
|
| model: An instance of the `YOLOv8` class containing the YOLOv8 model.
|
|
|
| Returns:
|
| None
|
|
|
| Raises:
|
| None
|
| """
|
| source_webcam = settings.WEBCAM_PATH
|
| is_display_tracker, tracker = display_tracker_options()
|
| if st.sidebar.button('Detect Objects'):
|
| try:
|
| vid_cap = cv2.VideoCapture(source_webcam)
|
| st_frame = st.empty()
|
| while (vid_cap.isOpened()):
|
| success, image = vid_cap.read()
|
| if success:
|
| _display_detected_frames(conf,
|
| model,
|
| st_frame,
|
| image,
|
| is_display_tracker,
|
| tracker,
|
| )
|
| else:
|
| vid_cap.release()
|
| break
|
| except Exception as e:
|
| st.sidebar.error("Error loading video: " + str(e))
|
|
|
|
|
| def play_stored_video(conf, model):
|
| """
|
| Plays a stored video file. Tracks and detects objects in real-time using the YOLOv8 object detection model.
|
|
|
| Parameters:
|
| conf: Confidence of YOLOv8 model.
|
| model: An instance of the `YOLOv8` class containing the YOLOv8 model.
|
|
|
| Returns:
|
| None
|
|
|
| Raises:
|
| None
|
| """
|
| source_vid = st.sidebar.selectbox(
|
| "Choose a video...", settings.VIDEOS_DICT.keys())
|
|
|
| is_display_tracker, tracker = display_tracker_options()
|
|
|
| with open(settings.VIDEOS_DICT.get(source_vid), 'rb') as video_file:
|
| video_bytes = video_file.read()
|
| if video_bytes:
|
| st.video(video_bytes)
|
|
|
| if st.sidebar.button('Detect Video Objects'):
|
| try:
|
| vid_cap = cv2.VideoCapture(
|
| str(settings.VIDEOS_DICT.get(source_vid)))
|
| st_frame = st.empty()
|
| while (vid_cap.isOpened()):
|
| success, image = vid_cap.read()
|
| if success:
|
| _display_detected_frames(conf,
|
| model,
|
| st_frame,
|
| image,
|
| is_display_tracker,
|
| tracker
|
| )
|
| else:
|
| vid_cap.release()
|
| break
|
| except Exception as e:
|
| st.sidebar.error("Error loading video: " + str(e))
|
|
|
