DanielNRU/plan_points

1import os2import cv23import numpy as np4import gradio as gr5from ultralytics import YOLO6 7def simplify_polygon(points, tol_ratio=0.02):8    """9    Упрощает полигон с помощью аппроксимации по Перу.10    11    Параметры:12        points (np.ndarray): Исходный набор точек полигона (Nx2).13        tol_ratio (float): Коэффициент допуска для аппроксимации.14        15    Возвращает:16        np.ndarray: Упрощённый набор точек полигона.17    """18    pts = points.reshape((-1, 1, 2)).astype(np.float32)19    perimeter = cv2.arcLength(pts, True)20    epsilon = tol_ratio * perimeter21    approx = cv2.approxPolyDP(pts, epsilon, True)22    return approx.reshape(-1, 2)23 24# Загружаем модель YOLO (файл model.pt должен находиться в корневой директории Space)25model_path = os.path.join(os.getcwd(), "model.pt")26model12 = YOLO(model_path)27 28def process_image(image):29    """30    Обрабатывает изображение:31      - Выполняет предсказание модели YOLO для детекции комнат.32      - Рисует упрощённые полигоны (с углами) для объектов класса 'room' (id = 2).33      - Формирует текст со списком координат углов.34      35    Параметры:36        image (np.ndarray): Исходное изображение.37        38    Возвращает:39        annotated_img (np.ndarray): Изображение с нанесёнными аннотациями.40        coord_text (str): Текст с координатами углов.41    """42    pred_results = list(model12.predict(source=image, imgsz=640))43    if len(pred_results) == 0:44        return image, "Ошибка: предсказания не получены."45    46    pred = pred_results[0]47    if pred.masks is None or pred.boxes is None:48        return image, "На изображении объекты не обнаружены."49    50    # Преобразуем исходное изображение в RGB для корректного отображения51    orig_img = pred.orig_img52    if orig_img.shape[2] == 3:53        original_rgb = cv2.cvtColor(orig_img.copy(), cv2.COLOR_BGR2RGB)54    else:55        original_rgb = orig_img.copy()56    annotated_img = original_rgb.copy()57    58    # Извлекаем полигоны и классы59    polygons = pred.masks.xy60    class_ids = pred.boxes.cls.cpu().numpy()61    62    coordinates_list = []63    for idx, polygon in enumerate(polygons):64        class_id = int(class_ids[idx])65        if class_id == 2:  # Обрабатываем только объекты класса "room"66            poly_np = np.array(polygon, dtype=np.float32)67            simplified_poly = simplify_polygon(poly_np, tol_ratio=0.005)68            coordinates_list.append({69                "object": idx + 1,70                "class": class_id,71                "coordinates": [(float(pt[0]), float(pt[1])) for pt in simplified_poly]72            })73            cv2.polylines(annotated_img, [simplified_poly.astype(np.int32)], isClosed=True, color=(0, 255, 0), thickness=2)74            for point in simplified_poly:75                cv2.circle(annotated_img, (int(point[0]), int(point[1])), 3, (255, 0, 0), -1)76    77    if coordinates_list:78        coord_text = ""79        for item in coordinates_list:80            coord_text += f"Координаты углов в комнате №: {item['object']}\n"81            for coord in item['coordinates']:82                coord_text += f"({coord[0]:.2f}, {coord[1]:.2f})\n"83            coord_text += "-" * 30 + "\n"84    else:85        coord_text = "На изображении не обнаружены комнаты."86    87    return annotated_img, coord_text88 89# Создаем Gradio-интерфейс90iface = gr.Interface(91    fn=process_image,92    inputs=gr.Image(type="numpy", label="Исходное изображение"),93    outputs=[94        gr.Image(type="numpy", label="Изображение с обнаруженными комнатами"),95        gr.Textbox(label="Координаты углов комнат", lines=10)96    ],97    title="Детекция комнат с использованием YOLO и Gradio",98    description="Загрузите изображение для детекции комнат и отображения углов."99)100 101iface.launch()