299 lines
11 KiB
Python
299 lines
11 KiB
Python
|
#!/usr/bin/env python3
|
|||
|
|
# -*- coding: utf-8 -*-
|
|||
|
|
"""
|
|||
|
|
Layout Detector — разделение страницы чертежа на зоны.
|
|||
|
|
|
|||
|
|
Зоны:
|
|||
|
|
- "drawing" — схемы, виды, разрезы (линии + текст, разрежено)
|
|||
|
|
- "table" — таблицы (плотные линии в сетке)
|
|||
|
|
- "title_block" — штамп (нижний правый угол или низ страницы)
|
|||
|
|
- "notes" — примечания, текстовые блоки
|
|||
|
|
- "legend" — легенда/условные обозначения
|
|||
|
|
|
|||
|
|
Алгоритм:
|
|||
|
|
1. Находит все линии на странице
|
|||
|
|
2. Находит прямоугольники = таблицы
|
|||
|
|
3. Анализирует плотность OCR текста
|
|||
|
|
4. Классифицирует регионы
|
|||
|
|
"""
|
|||
|
|
|
|||
|
|
import sys
|
|||
|
|
import json
|
|||
|
|
from pathlib import Path
|
|||
|
|
from typing import List, Dict, Tuple
|
|||
|
|
import cv2
|
|||
|
|
import numpy as np
|
|||
|
|
from PIL import Image
|
|||
|
|
|
|||
|
|
|
|||
|
|
def find_all_lines(img_gray: np.ndarray, min_length: int = 40):
|
|||
|
|
"""Находит все прямые линии (горизонтальные и вертикальные)."""
|
|||
|
|
_, binary = cv2.threshold(img_gray, 180, 255, cv2.THRESH_BINARY_INV)
|
|||
|
|
h, w = binary.shape
|
|||
|
|
lines = []
|
|||
|
|
|
|||
|
|
# Горизонтальные
|
|||
|
|
for y in range(h):
|
|||
|
|
row = binary[y, :]
|
|||
|
|
in_line = False
|
|||
|
|
start = 0
|
|||
|
|
for x in range(w):
|
|||
|
|
if row[x] > 128:
|
|||
|
|
if not in_line:
|
|||
|
|
in_line = True
|
|||
|
|
start = x
|
|||
|
|
else:
|
|||
|
|
if in_line:
|
|||
|
|
length = x - start
|
|||
|
|
if length >= min_length:
|
|||
|
|
lines.append(("h", start, y, x-1, y))
|
|||
|
|
in_line = False
|
|||
|
|
if in_line:
|
|||
|
|
length = w - start
|
|||
|
|
if length >= min_length:
|
|||
|
|
lines.append(("h", start, y, w-1, y))
|
|||
|
|
|
|||
|
|
# Вертикальные
|
|||
|
|
for x in range(w):
|
|||
|
|
col = binary[:, x]
|
|||
|
|
in_line = False
|
|||
|
|
start = 0
|
|||
|
|
for y in range(h):
|
|||
|
|
if col[y] > 128:
|
|||
|
|
if not in_line:
|
|||
|
|
in_line = True
|
|||
|
|
start = y
|
|||
|
|
else:
|
|||
|
|
if in_line:
|
|||
|
|
length = y - start
|
|||
|
|
if length >= min_length:
|
|||
|
|
lines.append(("v", x, start, x, y-1))
|
|||
|
|
in_line = False
|
|||
|
|
if in_line:
|
|||
|
|
length = h - start
|
|||
|
|
if length >= min_length:
|
|||
|
|
lines.append(("v", x, start, x, h-1))
|
|||
|
|
|
|||
|
|
return lines
|
|||
|
|
|
|||
|
|
|
|||
|
|
def find_rectangles(lines: List[Tuple], min_size: int = 100) -> List[Dict]:
|
|||
|
|
"""Находит прямоугольники, образованные пересечением линий."""
|
|||
|
|
horiz = [(l[1], l[2], l[3], l[4]) for l in lines if l[0] == "h"]
|
|||
|
|
vert = [(l[1], l[2], l[3], l[4]) for l in lines if l[0] == "v"]
|
|||
|
|
|
|||
|
|
# Группируем горизонтальные по Y
|
|||
|
|
from collections import defaultdict
|
|||
|
|
h_by_y = defaultdict(list)
|
|||
|
|
for x1, y1, x2, y2 in horiz:
|
|||
|
|
h_by_y[y1].append((x1, x2))
|
|||
|
|
|
|||
|
|
# Группируем вертикальные по X
|
|||
|
|
v_by_x = defaultdict(list)
|
|||
|
|
for x1, y1, x2, y2 in vert:
|
|||
|
|
v_by_x[x1].append((y1, y2))
|
|||
|
|
|
|||
|
|
rects = []
|
|||
|
|
# Ищем пары горизонтальных линий с общими вертикальными
|
|||
|
|
y_vals = sorted(h_by_y.keys())
|
|||
|
|
for i in range(len(y_vals)):
|
|||
|
|
for j in range(i+1, len(y_vals)):
|
|||
|
|
y_top = y_vals[i]
|
|||
|
|
y_bottom = y_vals[j]
|
|||
|
|
# Ищем общий X-интервал
|
|||
|
|
for x1_a, x2_a in h_by_y[y_top]:
|
|||
|
|
for x1_b, x2_b in h_by_y[y_bottom]:
|
|||
|
|
x_left = max(x1_a, x1_b)
|
|||
|
|
x_right = min(x2_a, x2_b)
|
|||
|
|
if x_right - x_left < min_size:
|
|||
|
|
continue
|
|||
|
|
# Проверяем, есть ли вертикальные линии на x_left и x_right
|
|||
|
|
has_left = any(y_top <= y_bottom and not (y2 < y_top or y1 > y_bottom)
|
|||
|
|
for y1, y2 in v_by_x.get(x_left, []))
|
|||
|
|
has_right = any(y_top <= y_bottom and not (y2 < y_top or y1 > y_bottom)
|
|||
|
|
for y1, y2 in v_by_x.get(x_right, []))
|
|||
|
|
if has_left and has_right:
|
|||
|
|
rects.append({
|
|||
|
|
"x": x_left, "y": y_top,
|
|||
|
|
"w": x_right - x_left, "h": y_bottom - y_top
|
|||
|
|
})
|
|||
|
|
|
|||
|
|
# Фильтруем вложенные прямоугольники (оставляем только внешние)
|
|||
|
|
filtered = []
|
|||
|
|
for r in rects:
|
|||
|
|
is_inner = False
|
|||
|
|
for other in rects:
|
|||
|
|
if r is other:
|
|||
|
|
continue
|
|||
|
|
if (r["x"] > other["x"] and r["y"] > other["y"] and
|
|||
|
|
r["x"] + r["w"] < other["x"] + other["w"] and
|
|||
|
|
r["y"] + r["h"] < other["y"] + other["h"]):
|
|||
|
|
is_inner = True
|
|||
|
|
break
|
|||
|
|
if not is_inner:
|
|||
|
|
filtered.append(r)
|
|||
|
|
|
|||
|
|
return filtered
|
|||
|
|
|
|||
|
|
|
|||
|
|
def classify_regions(rects: List[Dict], ocr_lines: List[Dict], img_w: int, img_h: int) -> List[Dict]:
|
|||
|
|
"""Классифицирует регионы страницы."""
|
|||
|
|
regions = []
|
|||
|
|
|
|||
|
|
# 1. Таблицы = большие прямоугольники с высокой плотностью линий
|
|||
|
|
for r in rects:
|
|||
|
|
area = r["w"] * r["h"]
|
|||
|
|
# Считаем OCR строки внутри
|
|||
|
|
texts_in = [t for t in ocr_lines
|
|||
|
|
if r["x"] <= t["cx"] <= r["x"] + r["w"]
|
|||
|
|
and r["y"] <= t["cy"] <= r["y"] + r["h"]]
|
|||
|
|
density = len(texts_in) / (area / 1000000) # текстов на мегапиксель
|
|||
|
|
|
|||
|
|
if density > 20: # высокая плотность = таблица
|
|||
|
|
regions.append({
|
|||
|
|
"type": "table",
|
|||
|
|
"bbox": [r["x"], r["y"], r["x"]+r["w"], r["y"]+r["h"]],
|
|||
|
|
"density": density,
|
|||
|
|
"text_count": len(texts_in)
|
|||
|
|
})
|
|||
|
|
|
|||
|
|
# 2. Определяем чертежи = области с линиями и текстом, но без плотной сетки
|
|||
|
|
# Для простоты: левая половина, не покрытая таблицами
|
|||
|
|
# Найдём ограничивающий bbox для всех "чертёжных" текстов
|
|||
|
|
drawing_texts = [t for t in ocr_lines if t["cy"] < img_h * 0.75 and t["cx"] < img_w * 0.6]
|
|||
|
|
if drawing_texts:
|
|||
|
|
xs = [t["cx"] for t in drawing_texts]
|
|||
|
|
ys = [t["cy"] for t in drawing_texts]
|
|||
|
|
# Расширяем на 200px
|
|||
|
|
dx = [t["cx"] - t["x1"] for t in drawing_texts if "x1" in t]
|
|||
|
|
max_w = max(dx) if dx else 100
|
|||
|
|
regions.append({
|
|||
|
|
"type": "drawing",
|
|||
|
|
"bbox": [max(0, min(xs)-max_w), max(0, min(ys)-100),
|
|||
|
|
min(img_w, max(xs)+max_w), min(img_h, max(ys)+100)],
|
|||
|
|
"text_count": len(drawing_texts)
|
|||
|
|
})
|
|||
|
|
|
|||
|
|
# 3. Штамп = низ страницы, мелкий текст
|
|||
|
|
title_texts = [t for t in ocr_lines if t["cy"] > img_h * 0.85]
|
|||
|
|
if title_texts:
|
|||
|
|
xs = [t["cx"] for t in title_texts]
|
|||
|
|
ys = [t["cy"] for t in title_texts]
|
|||
|
|
regions.append({
|
|||
|
|
"type": "title_block",
|
|||
|
|
"bbox": [min(xs)-50, min(ys)-50, max(xs)+50, max(ys)+50],
|
|||
|
|
"text_count": len(title_texts)
|
|||
|
|
})
|
|||
|
|
|
|||
|
|
# 4. Примечания = текстовые блоки
|
|||
|
|
note_keywords = ["примечание", "общие указания", "границы", "размеры"]
|
|||
|
|
note_texts = [t for t in ocr_lines
|
|||
|
|
if any(kw in t["text"].lower() for kw in note_keywords)]
|
|||
|
|
if note_texts:
|
|||
|
|
xs = [t["cx"] for t in note_texts]
|
|||
|
|
ys = [t["cy"] for t in note_texts]
|
|||
|
|
regions.append({
|
|||
|
|
"type": "notes",
|
|||
|
|
"bbox": [min(xs)-100, min(ys)-100, max(xs)+100, max(ys)+100],
|
|||
|
|
"text_count": len(note_texts)
|
|||
|
|
})
|
|||
|
|
|
|||
|
|
return regions
|
|||
|
|
|
|||
|
|
|
|||
|
|
def detect_layout(png_path: Path, ocr_path: Path) -> Dict:
|
|||
|
|
"""Основная функция layout detection."""
|
|||
|
|
img = cv2.imread(str(png_path), cv2.IMREAD_GRAYSCALE)
|
|||
|
|
h, w = img.shape[:2]
|
|||
|
|
|
|||
|
|
# Загрузить OCR
|
|||
|
|
ocr = json.loads(ocr_path.read_text(encoding="utf-8"))
|
|||
|
|
|
|||
|
|
# Собрать все OCR lines с координатами
|
|||
|
|
all_texts = []
|
|||
|
|
for page in ocr.get("pages", []):
|
|||
|
|
for line in page.get("ocr_lines", []):
|
|||
|
|
bbox = line.get("bbox", [])
|
|||
|
|
if not bbox:
|
|||
|
|
continue
|
|||
|
|
if isinstance(bbox[0], list):
|
|||
|
|
xs = [p[0] for p in bbox]
|
|||
|
|
ys = [p[1] for p in bbox]
|
|||
|
|
else:
|
|||
|
|
xs = [bbox[0], bbox[2]]
|
|||
|
|
ys = [bbox[1], bbox[3]]
|
|||
|
|
all_texts.append({
|
|||
|
|
"text": line["text"],
|
|||
|
|
"cx": sum(xs)/len(xs),
|
|||
|
|
"cy": sum(ys)/len(ys),
|
|||
|
|
"x1": min(xs), "y1": min(ys),
|
|||
|
|
"x2": max(xs), "y2": max(ys),
|
|||
|
|
"bbox": bbox
|
|||
|
|
})
|
|||
|
|
|
|||
|
|
# Найти линии
|
|||
|
|
lines = find_all_lines(img)
|
|||
|
|
print(f"[INFO] Найдено {len(lines)} линий")
|
|||
|
|
|
|||
|
|
# Найти прямоугольники
|
|||
|
|
rects = find_rectangles(lines)
|
|||
|
|
print(f"[INFO] Найдено {len(rects)} прямоугольников")
|
|||
|
|
|
|||
|
|
# Классифицировать
|
|||
|
|
regions = classify_regions(rects, all_texts, w, h)
|
|||
|
|
print(f"[INFO] Классифицировано {len(regions)} регионов")
|
|||
|
|
for r in regions:
|
|||
|
|
print(f" {r['type']}: bbox={r['bbox']}, texts={r.get('text_count', 0)}")
|
|||
|
|
|
|||
|
|
return {
|
|||
|
|
"image_size": [w, h],
|
|||
|
|
"regions": regions,
|
|||
|
|
"rectangles": rects,
|
|||
|
|
"line_count": len(lines)
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
|
|||
|
|
def visualize_layout(png_path: Path, layout: Dict, out_path: Path):
|
|||
|
|
"""Рисует зоны на изображении."""
|
|||
|
|
img = Image.open(png_path)
|
|||
|
|
draw = ImageDraw.Draw(img)
|
|||
|
|
colors = {
|
|||
|
|
"table": "blue",
|
|||
|
|
"drawing": "green",
|
|||
|
|
"title_block": "purple",
|
|||
|
|
"notes": "orange"
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
for region in layout["regions"]:
|
|||
|
|
x1, y1, x2, y2 = region["bbox"]
|
|||
|
|
color = colors.get(region["type"], "red")
|
|||
|
|
draw.rectangle([x1, y1, x2, y2], outline=color, width=4)
|
|||
|
|
draw.text((x1+5, y1+5), region["type"], fill=color)
|
|||
|
|
|
|||
|
|
img.save(out_path)
|
|||
|
|
print(f"[OK] Layout visualization: {out_path}")
|
|||
|
|
|
|||
|
|
|
|||
|
|
def main():
|
|||
|
|
if len(sys.argv) < 3:
|
|||
|
|
print("Usage: python layout_detector.py <png> <ocr_json>")
|
|||
|
|
sys.exit(1)
|
|||
|
|
|
|||
|
|
png = Path(sys.argv[1])
|
|||
|
|
ocr = Path(sys.argv[2])
|
|||
|
|
out_json = png.parent / "layout.json"
|
|||
|
|
out_png = png.parent / f"{png.stem}_layout.png"
|
|||
|
|
|
|||
|
|
layout = detect_layout(png, ocr)
|
|||
|
|
|
|||
|
|
with open(out_json, "w", encoding="utf-8") as f:
|
|||
|
|
json.dump(layout, f, ensure_ascii=False, indent=2)
|
|||
|
|
print(f"[OK] Layout JSON: {out_json}")
|
|||
|
|
|
|||
|
|
visualize_layout(png, layout, out_png)
|
|||
|
|
|
|||
|
|
|
|||
|
|
if __name__ == "__main__":
|
|||
|
|
from PIL import ImageDraw
|
|||
|
|
main()
|