Update app.py

app.py

@@ -1,4 +1,48 @@
 import gradio as gr
+import cv2
+import numpy as np
+from ultralytics import YOLO
+
+# Загрузка модели (путь должен быть относительным или абсолютным)
+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)
@@ -7,8 +51,10 @@ demo = gr.Interface(
     fn=coordinates,
     inputs=[gr.Image(type="numpy")],
     outputs=[gr.Image(label="Результат детекции"), gr.Textbox(label="Обнаруженные объекты", lines=10)],
-    title="Детектор дверей и окон"
+    title="Детекция дверей и окон"
 )
 
 if __name__ == "__main__":
     demo.launch()
+
+