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)