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()