live-cam / app.py
vlasdadsda's picture
Update app.py
946cb95 verified
import gradio as gr
import cv2
from ultralytics import YOLO
# Загружаем модель YOLO
model_path = "sss.pt" # Замените на путь к вашей модели
model = YOLO(model_path)
def process_frame(frame):
# Выполняем обнаружение объектов
results = model(frame, stream=True) # Используем stream=True для повышения скорости
# Рисуем боксы и метки на кадре
for result in results:
for box in result.boxes.data.tolist():
x1, y1, x2, y2, confidence, class_id = box
x1, y1, x2, y2 = int(x1), int(y1), int(x2), int(y2)
label = model.names[int(class_id)] # Убираем уверенность, оставляем только имя объекта
# Рисуем прямоугольник вокруг объекта
cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 255, 0), 2)
# Добавляем текст
cv2.putText(frame, label, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)
return frame
def capture_video():
# Открываем видеопоток с камеры
cap = cv2.VideoCapture(0)
if not cap.isOpened():
print("Не удалось открыть камеру")
return
while True:
ret, frame = cap.read()
if not ret:
break
# Обрабатываем кадр
processed_frame = process_frame(frame)
# Gradio ожидает изображение в формате NumPy, поэтому передаем его так
yield processed_frame
cap.release()
# Создаем интерфейс Gradio
iface = gr.Interface(fn=capture_video,
inputs=[],
outputs=gr.Image(type="numpy", label="Object Detection"), # Ожидаем NumPy массив
live=True)
# Запуск интерфейса
iface.launch()