app.py

import os
import cv2
import numpy as np
import gradio as gr
from ultralytics import YOLO

def fit_circle(points):
    """
    Аппроксимирует дугу окружностью через cv2.minEnclosingCircle.
    Возвращает центр и радиус.
    """
    pts = np.array(points, dtype=np.float32)
    center, radius = cv2.minEnclosingCircle(pts)
    return tuple(center), radius

def simplify_polygon(points, rdp_tol=0.005, curvature_threshold=15, arc_resolution=20):
    """
    Упрощает полигон с помощью алгоритма RDP и аппроксимирует дуги.
    """
    points = np.array(points, dtype=np.float32)
    epsilon = rdp_tol * cv2.arcLength(points, True)
    simplified = cv2.approxPolyDP(points, epsilon=epsilon, closed=True).reshape(-1, 2)

    optimized = []
    n = len(simplified)
    i = 0
    while i < n:
        prev = simplified[i - 1]
        curr = simplified[i]
        nxt = simplified[(i + 1) % n]

        v1 = curr - prev
        v2 = nxt - curr
        norm_v1 = np.linalg.norm(v1)
        norm_v2 = np.linalg.norm(v2)
        angle_deg = 0 if norm_v1 < 1e-6 or norm_v2 < 1e-6 else np.degrees(
            np.arccos(np.clip(np.dot(v1, v2) / (norm_v1 * norm_v2), -1.0, 1.0))
        )

        if abs(angle_deg - 180) > curvature_threshold:
            optimized.append(curr.tolist())
            i += 1
        else:
            arc_points = [curr]
            j = i + 1
            while j < n:
                prev_j = simplified[j - 1]
                curr_j = simplified[j]
                nxt_j = simplified[(j + 1) % n]
                v1_j = curr_j - prev_j
                v2_j = nxt_j - curr_j
                norm_v1_j = np.linalg.norm(v1_j)
                norm_v2_j = np.linalg.norm(v2_j)
                angle_j = 0 if norm_v1_j < 1e-6 or norm_v2_j < 1e-6 else np.degrees(
                    np.arccos(np.clip(np.dot(v1_j, v2_j) / (norm_v1_j * norm_v2_j), -1.0, 1.0))
                )
                if abs(angle_j - 180) > curvature_threshold:
                    break
                arc_points.append(curr_j)
                j += 1

            if len(arc_points) >= 3:
                center, radius = fit_circle(arc_points)
                start_angle = np.arctan2(arc_points[0][1] - center[1], arc_points[0][0] - center[0])
                end_angle = np.arctan2(arc_points[-1][1] - center[1], arc_points[-1][0] - center[0])
                if end_angle < start_angle:
                    end_angle += 2 * np.pi
                interp_angles = np.linspace(start_angle, end_angle, arc_resolution)
                arc_interp = [
                    [int(center[0] + radius * np.cos(a)), int(center[1] + radius * np.sin(a))]
                    for a in interp_angles
                ]
                optimized.extend(arc_interp)
            else:
                optimized.extend([pt.tolist() for pt in arc_points])
            i = j
    return np.array(optimized, dtype=np.int32)

def filter_straight_segments(points, tolerance=1.0):
    """
    Фильтрует промежуточные точки на прямых участках.
    Удаляет точку, если площадь треугольника, образованного соседними точками, меньше tolerance.
    """
    if points.shape[0] < 3:
        return points
    filtered = [points[0]]
    for i in range(1, len(points) - 1):
        p_prev = filtered[-1]
        p_curr = points[i]
        p_next = points[i + 1]
        area = 0.5 * np.abs(np.cross(p_next - p_prev, p_curr - p_prev))
        if area > tolerance:
            filtered.append(p_curr)
    filtered.append(points[-1])
    return np.array(filtered)

def straighten_points(points, angle_threshold=15, iterations=3):
    """
    Корректирует точки, если угол между отрезками, соединяющими вершину с соседями,
    близок к 180° (почти прямая линия), посредством проекции на прямую.
    """
    n = len(points)
    if n < 3:
        return points
    new_points = points.astype(np.float32).copy()
    for _ in range(iterations):
        for i in range(n):
            prev = new_points[i - 1]
            curr = new_points[i]
            nxt = new_points[(i + 1) % n]
            v1 = prev - curr
            v2 = nxt - curr
            norm1 = np.linalg.norm(v1)
            norm2 = np.linalg.norm(v2)
            if norm1 < 1e-6 or norm2 < 1e-6:
                continue
            cos_angle = np.dot(v1, v2) / (norm1 * norm2)
            angle = np.degrees(np.arccos(np.clip(cos_angle, -1.0, 1.0)))
            if (180 - angle) < angle_threshold:
                line_vec = nxt - prev
                if np.linalg.norm(line_vec) < 1e-6:
                    continue
                t = np.dot(curr - prev, line_vec) / np.dot(line_vec, line_vec)
                new_points[i] = prev + t * line_vec
    return new_points.astype(np.int32)

def align_points_x_rel(points, proportion=0.02):
    """
    Выравнивает соседние точки по оси X, если разница меньше (proportion * ширина полигона).
    """
    new_points = points.copy().astype(np.int32)
    xs = new_points[:, 0]
    width = xs.max() - xs.min()
    threshold = proportion * width
    for i in range(1, len(new_points)):
        if abs(new_points[i][0] - new_points[i-1][0]) < threshold:
            new_points[i][0] = new_points[i-1][0]
    return new_points

def align_points_y_rel(points, proportion=0.02):
    """
    Выравнивает соседние точки по оси Y, если разница меньше (proportion * высота полигона).
    """
    new_points = points.copy().astype(np.int32)
    ys = new_points[:, 1]
    height = ys.max() - ys.min()
    threshold = proportion * height
    for i in range(1, len(new_points)):
        if abs(new_points[i][1] - new_points[i-1][1]) < threshold:
            new_points[i][1] = new_points[i-1][1]
    return new_points

def align_polygon_edges(points, angle_threshold=10):
    """
    Выравнивает сегменты замкнутого полигона:
    точки между угловыми проецируются на прямую, соединяющую крайние точки группы.
    """
    n = len(points)
    if n < 3:
        return points

    angles = []
    for i in range(n):
        prev = points[i - 1]
        curr = points[i]
        nxt = points[(i + 1) % n]
        v1 = prev - curr
        v2 = nxt - curr
        norm1 = np.linalg.norm(v1)
        norm2 = np.linalg.norm(v2)
        angle = 180 if norm1 < 1e-6 or norm2 < 1e-6 else np.degrees(
            np.arccos(np.clip(np.dot(v1, v2) / (norm1 * norm2), -1.0, 1.0))
        )
        angles.append(angle)

    corner_idxs = [i for i, a in enumerate(angles) if (180 - a) > angle_threshold]
    if not corner_idxs:
        return points

    if corner_idxs[0] != 0:
        corner_idxs = [0] + corner_idxs
    if corner_idxs[-1] != (n - 1):
        corner_idxs.append(n - 1)

    new_points = points.astype(np.float32).copy()
    for i in range(len(corner_idxs) - 1):
        start_idx = corner_idxs[i]
        end_idx = corner_idxs[i + 1]
        start_pt = new_points[start_idx]
        end_pt = new_points[end_idx]
        d = end_pt - start_pt
        norm_d_sq = np.dot(d, d)
        if norm_d_sq < 1e-6:
            continue
        for j in range(start_idx + 1, end_idx):
            t = np.dot(new_points[j] - start_pt, d) / norm_d_sq
            new_points[j] = start_pt + t * d
    return new_points.astype(np.int32)

def process_image(image, filtration):
    """
    Обрабатывает изображение:
      - Детектирует объекты моделью YOLO.
      - Для объектов класса 'room' (id == 2) извлекается полигон.
      - При включённой фильтрации применяется цепочка обработки:
          • Упрощение (RDP + аппроксимация дуг)
          • Фильтрация промежуточных точек
          • Выравнивание точек (проекция почти прямых вершин)
          • Относительное выравнивание по осям
          • Выравнивание по сторонам (проекция точек сегментов на прямую)
          • Финальная аппроксимация до 4-х точек для прямоугольных форм.
      - Рисуется полигон и подпись с названием комнаты.
      - Формируется текст с координатами опорных точек.
    """
    pred_results = list(model.predict(source=image, imgsz=640))
    if not pred_results:
        return image, "Ошибка: предсказания не получены."

    pred = pred_results[0]
    if pred.masks is None or pred.boxes is None:
        return image, "На изображении объекты не обнаружены."

    image_rgb = cv2.cvtColor(pred.orig_img.copy(), cv2.COLOR_BGR2RGB)
    polygons = pred.masks.xy
    class_ids = pred.boxes.cls.cpu().numpy()

    img_h, img_w = image_rgb.shape[:2]
    font_scale = (img_w + img_h) / 2000.0
    thickness = max(int(round(font_scale * 2)), 1)

    coordinates_list = []
    for idx, polygon in enumerate(polygons):
        class_id = int(class_ids[idx])
        if class_id == 2:
            poly_np = np.array(polygon, dtype=np.float32)
            if filtration:
                processed_poly = simplify_polygon(poly_np, rdp_tol=0.005, curvature_threshold=15, arc_resolution=20)
                processed_poly = filter_straight_segments(processed_poly, tolerance=1.0)
                processed_poly = straighten_points(processed_poly, angle_threshold=15, iterations=3)
                processed_poly = align_points_x_rel(processed_poly, proportion=0.02)
                processed_poly = align_points_y_rel(processed_poly, proportion=0.02)
                processed_poly = align_polygon_edges(processed_poly, angle_threshold=10)
                epsilon_final = 0.01 * cv2.arcLength(processed_poly, True)
                approx_final = cv2.approxPolyDP(processed_poly, epsilon=epsilon_final, closed=True)
                if len(approx_final) == 4:
                    processed_poly = approx_final.reshape(-1, 2)
            else:
                processed_poly = poly_np.astype(np.int32)

            coordinates_list.append({
                "object": idx + 1,
                "class": class_id,
                "coordinates": [(float(pt[0]), float(pt[1])) for pt in processed_poly]
            })
            cv2.polylines(image_rgb, [processed_poly.astype(np.int32)], isClosed=True, color=(0, 255, 0), thickness=2)
            for point in processed_poly:
                cv2.circle(image_rgb, (int(point[0]), int(point[1])), 3, (255, 0, 0), -1)

            x, y, _, _ = cv2.boundingRect(processed_poly)
            label = f"Room #{idx + 1}"
            pos = (x, y - 5 if y - 5 > 0 else y)
            cv2.putText(image_rgb, label, pos, cv2.FONT_HERSHEY_SIMPLEX, font_scale, (0, 0, 255), thickness, cv2.LINE_AA)

    if coordinates_list:
        coord_text = ""
        for item in coordinates_list:
            coord_text += f"Комната {item['object']}:\nКоординаты точек:\n"
            for coord in item['coordinates']:
                coord_text += f"({coord[0]:.2f}, {coord[1]:.2f})\n"
            coord_text += "-" * 30 + "\n"
    else:
        coord_text = "На изображении объекты 'room' не обнаружены."

    return image_rgb, coord_text

# Загружаем модель YOLO (файл model.pt должен находиться в корневой директории)
model_path = os.path.join(os.getcwd(), "model.pt")
model = YOLO(model_path)

iface = gr.Interface(
    fn=process_image,
    inputs=[
        gr.Image(type="numpy", label="Исходное изображение"),
        gr.Checkbox(label="Фильтрация", value=True)
    ],
    outputs=[
        gr.Image(type="numpy", label="Изображение с обнаруженными комнатами"),
        gr.Textbox(label="Координаты опорных точек", lines=10)
    ],
    title="Детекция комнат с использованием YOLO и Gradio",
    description="Загрузите изображение для детекции комнат."
)

iface.launch()