Spaces:
Sleeping
Sleeping
File size: 2,795 Bytes
add170b c166f7d b106e83 c166f7d ec000c4 c166f7d add170b 2611860 c166f7d 2611860 add170b 3751ddb c166f7d | 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 | import gradio as gr
import cv2
import numpy as np
from ultralytics import YOLO
import os
os.environ["YOLO_CONFIG_DIR"] = "/tmp/ultralytics_config"
os.makedirs(os.environ["YOLO_CONFIG_DIR"], exist_ok=True)
# Загрузка модели (путь должен быть относительным или абсолютным)
model = YOLO("models/best.pt")
def draw(image_for_drawing):
#detections = []
results = model.predict(image_for_drawing, conf=0.5)
#Перебираем результаты
for result in results:
# Получаем bounding boxes в формате [x1, y1, x2, y2, conf, class_id]
boxes = result.boxes.xyxy.cpu().numpy() # Координаты в пикселях
confidences = result.boxes.conf.cpu().numpy() # Уверенность
class_ids = result.boxes.cls.cpu().numpy().astype(int) # ID классов
door_boxes = boxes[class_ids == 1] # Массив координат дверей
windows_boxes = boxes[class_ids == 4]
door_confs = confidences[class_ids == 1] # Уверенность для дверей
windows_confs = confidences[class_ids == 4] # Уверенность для окон
detections = []
# Выводим координаты и добавляем надписи
for box, conf in zip(door_boxes, door_confs):
x1, y1, x2, y2 = box
detections.append(f"Дверь: x1={x1:.0f}, y1={y1:.0f}, x2={x2:.0f}, y2={y2:.0f}, confidence={conf:.2f}")
cv2.rectangle(image_for_drawing, (int(x1), int(y1)), (int(x2), int(y2)), (255, 0, 0), 2)
# Добавляем текст
text = f"Door: {conf:.2f}"
cv2.putText(image_for_drawing, text, (int(x1), int(y1) - 10),
cv2.FONT_HERSHEY_SIMPLEX, 0.4, (255, 0, 0), 1)
for box, conf in zip(windows_boxes, windows_confs):
x1, y1, x2, y2 = box
detections.append(f"Окно: x1={x1:.0f}, y1={y1:.0f}, x2={x2:.0f}, y2={y2:.0f}, confidence={conf:.2f}")
cv2.rectangle(image_for_drawing, (int(x1), int(y1)), (int(x2), int(y2)), (0, 255, 0), 2)
# Добавляем текст
text = f"window: {conf:.2f}"
cv2.putText(image_for_drawing, text, (int(x1), int(y1) - 10),
cv2.FONT_HERSHEY_SIMPLEX, 0.4, (0, 255, 0), 1)
#return image_for_drawing
return image_for_drawing, "\n".join(detections)
def coordinates(image):
return draw(image)
demo = gr.Interface(
fn=coordinates,
inputs=[gr.Image(type="numpy")],
outputs=[gr.Image(label="Результат детекции"), gr.Textbox(label="Обнаруженные объекты", lines=10)],
title="Детекция дверей и окон"
)
if __name__ == "__main__":
demo.launch(share=True)
|