File size: 3,764 Bytes
3f67d40 |
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 |
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
)
@st.cache_resource
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)}") |