Spaces:
Runtime error
Runtime error
from ultralytics import YOLO | |
import streamlit as st | |
import cv2 | |
from PIL import Image | |
import tempfile | |
def _display_detected_frames(conf, model, st_frame, image): | |
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 load_model(model_path): | |
model = YOLO(model_path) | |
return model | |
def infer_uploaded_image(conf, model): | |
source_img = st.sidebar.file_uploader( | |
label="Выберите изображение...", | |
type=("jpg", "jpeg", "png", 'bmp', 'webp') | |
) | |
col1, col2 = st.columns(2) | |
with col1: | |
if source_img: | |
uploaded_image = Image.open(source_img) | |
st.image( | |
image=source_img, | |
caption="Загруженное изображение", | |
use_column_width=True | |
) | |
if source_img: | |
if st.button("Старт"): | |
with st.spinner("Запускается..."): | |
res = model.predict(uploaded_image, | |
conf=conf) | |
boxes = res[0].boxes | |
res_plotted = res[0].plot()[:, :, ::-1] | |
with col2: | |
st.image(res_plotted, | |
caption="Обработанное изображение", | |
use_column_width=True) | |
try: | |
with st.expander("Результаты обнаружения"): | |
for box in boxes: | |
st.write(box.xywh) | |
except Exception as ex: | |
st.write("Изображение не загружено") | |
st.write(ex) | |
def infer_uploaded_video(conf, model): | |
source_video = st.sidebar.file_uploader( | |
label="Выберите видео..." | |
) | |
if source_video: | |
st.video(source_video) | |
if source_video: | |
if st.button("Старт"): | |
with st.spinner("Запускается..."): | |
try: | |
tfile = tempfile.NamedTemporaryFile() | |
tfile.write(source_video.read()) | |
vid_cap = cv2.VideoCapture( | |
tfile.name) | |
st_frame = st.empty() | |
while (vid_cap.isOpened()): | |
success, image = vid_cap.read() | |
if success: | |
_display_detected_frames(conf, | |
model, | |
st_frame, | |
image | |
) | |
else: | |
vid_cap.release() | |
break | |
except Exception as e: | |
st.error(f"Ошибка загрузки видео: {e}") | |
def infer_uploaded_webcam(conf, model): | |
try: | |
flag = st.button( | |
label="Остановить" | |
) | |
vid_cap = cv2.VideoCapture(0) | |
st_frame = st.empty() | |
while not flag: | |
success, image = vid_cap.read() | |
if success: | |
_display_detected_frames( | |
conf, | |
model, | |
st_frame, | |
image | |
) | |
else: | |
vid_cap.release() | |
break | |
except Exception as e: | |
st.error(f"Ошибка загрузки видео: {str(e)}") |