Spaces:
Sleeping
Sleeping
# Python In-built packages | |
from pathlib import Path | |
import PIL | |
import time | |
# External packages | |
import streamlit as st | |
import cv2 | |
# Local Modules | |
import settings | |
import helper | |
# Setting page layout | |
st.set_page_config( | |
page_title="Object Detection using YOLOv9", | |
page_icon="🤖", | |
layout="wide", | |
initial_sidebar_state="expanded" | |
) | |
# Main page heading | |
st.title("Object Detection using YOLOv9") | |
# Sidebar | |
st.sidebar.header("ML Model Config") | |
# Model Options | |
model_type = st.sidebar.radio( | |
"Select Task", ['Detection', 'Segmentation']) | |
confidence = float(st.sidebar.slider( | |
"Select Model Confidence", 25, 100, 40)) / 100 | |
# Selecting Detection Or Segmentation | |
if model_type == 'Detection': | |
model_path = Path(settings.DETECTION_MODEL) | |
elif model_type == 'Segmentation': | |
model_path = Path(settings.SEGMENTATION_MODEL) | |
# Load Pre-trained ML Model | |
try: | |
model = helper.load_model(model_path) | |
except Exception as ex: | |
st.error(f"Unable to load model. Check the specified path: {model_path}") | |
st.error(ex) | |
st.sidebar.header("Image/Video Config") | |
source_radio = st.sidebar.radio( | |
"Select Source", settings.SOURCES_LIST) | |
source_img = None | |
# If image is selected | |
if source_radio == settings.IMAGE: | |
source_img = st.sidebar.file_uploader( | |
"Choose an image...", type=("jpg", "jpeg", "png", 'bmp', 'webp')) | |
col1, col2 = st.columns(2) | |
with col1: | |
try: | |
if source_img is None: | |
default_image_path = str(settings.DEFAULT_IMAGE) | |
default_image = PIL.Image.open(default_image_path) | |
st.image(default_image_path, caption="Default Image", | |
use_column_width=True) | |
else: | |
uploaded_image = PIL.Image.open(source_img) | |
st.image(source_img, caption="Uploaded Image", | |
use_column_width=True) | |
except Exception as ex: | |
st.error("Error occurred while opening the image.") | |
st.error(ex) | |
with col2: | |
if source_img is None: | |
default_detected_image_path = str(settings.DEFAULT_DETECT_IMAGE) | |
default_detected_image = PIL.Image.open( | |
default_detected_image_path) | |
st.image(default_detected_image_path, caption='Detected Image', | |
use_column_width=True) | |
else: | |
if st.sidebar.button('Detect Objects'): | |
res = model.predict(uploaded_image, | |
conf=confidence | |
) | |
boxes = res[0].boxes | |
res_plotted = res[0].plot()[:, :, ::-1] | |
st.image(res_plotted, caption='Detected Image', | |
use_column_width=True) | |
try: | |
with st.expander("Detection Results"): | |
for box in boxes: | |
st.write(box.data) | |
except Exception as ex: | |
# st.write(ex) | |
st.write("No image is uploaded yet!") | |
elif source_radio == settings.VIDEO: | |
helper.play_stored_video(confidence, model) | |
elif source_radio == settings.YOUTUBE: | |
helper.play_youtube_video(confidence, model) | |
elif source_radio == settings.WEBCAM: # Adding webcam option | |
st.sidebar.markdown("**Webcam**") | |
webcam = cv2.VideoCapture(1) | |
if not webcam.isOpened(): | |
st.error("Unable to open webcam. Please make sure it's connected.") | |
else: | |
webcam_status = st.sidebar.empty() | |
start_stop_button = st.sidebar.button("Start Webcam") | |
if start_stop_button: | |
webcam_started = True | |
webcam_status.text("Webcam Started!") | |
while webcam_started: | |
ret, frame = webcam.read() | |
if not ret: | |
st.error("Failed to capture frame from webcam") | |
break | |
frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) | |
# Perform object detection on the frame | |
res = model.predict(frame, conf=confidence) | |
boxes = res[0].boxes | |
frame = res[0].plot()[:, :, ::-1] | |
# Display the frame | |
st.image(frame, channels="RGB", use_column_width=True) | |
stop_button = st.sidebar.button("Stop Webcam" + str(hash(time.time()))) # Unique key | |
if stop_button: | |
webcam_started = False | |
webcam_status.text("Webcam Stopped!") | |
webcam.release() | |