Spaces:
Sleeping
Sleeping
| 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) | |