|
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))
|
|
|
