Spaces:
Sleeping
Sleeping
File size: 4,552 Bytes
22b2bcd |
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 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 |
# 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()
|