b5f7c6327e
- tiling_ocr.py: split large drawings into overlapping tiles for better small-text recognition - preprocess_for_ocr.py: CLAHE + unsharp mask for enhancing blueprint contrast - visualize_dimensions.py: draw bounding boxes around detected dimension numbers - compare_ocr.py: side-by-side visualization of normal vs tiling OCR results - dimension_extractor.py: line-based dimension detection with pixel verification - ocr_qwen.py: Alibaba Cloud qwen-vl-ocr client with resize and regex fallback parser - test_qwen_ocr.py: standalone test for qwen OCR - process_any_pdf.py: add --use-tiling flag to switch between normal and tiling OCR
194 lines
7.7 KiB
Python
194 lines
7.7 KiB
Python
#!/usr/bin/env python3
|
||
# -*- coding: utf-8 -*-
|
||
"""
|
||
Локальный детектор размеров на чертеже.
|
||
|
||
Подход:
|
||
1. Находим линии на PNG (Canny + HoughLinesP) — только горизонтальные/вертикальные
|
||
2. Загружаем OCR результаты, фильтруем только числа (regex ^\d+([,.]\d+)?$)
|
||
3. Для каждого числа проверяем: есть ли линия в радиусе 60px?
|
||
4. Если да — считаем это размером
|
||
5. Визуализируем результат
|
||
|
||
Результат: dimensions.json + *_dims_detected.png
|
||
"""
|
||
|
||
import sys
|
||
import json
|
||
import math
|
||
import re
|
||
from pathlib import Path
|
||
from typing import List, Dict, Tuple
|
||
import cv2
|
||
import numpy as np
|
||
from PIL import Image, ImageDraw
|
||
|
||
|
||
def find_numbers_with_context(ocr_path: Path, png_path: Path) -> Tuple[List[Dict], List[Dict]]:
|
||
"""
|
||
Находит размеры через анализ контекста:
|
||
1. Берём все числа из OCR
|
||
2. Ищем "соседей" на той же горизонтали/вертикали (размерные цепочки)
|
||
3. Исключаем числа из таблиц (по bbox: справа на странице)
|
||
4. Проверяем пиксели между числами: есть ли линия?
|
||
"""
|
||
print(f"[INFO] Обработка {png_path.name}...")
|
||
|
||
ocr = json.loads(ocr_path.read_text(encoding="utf-8"))
|
||
img = cv2.imread(str(png_path), cv2.IMREAD_GRAYSCALE)
|
||
h, w = img.shape[:2]
|
||
|
||
numbers = []
|
||
for page in ocr.get("pages", []):
|
||
for line_data in page.get("ocr_lines", []):
|
||
txt = line_data["text"].strip()
|
||
if not re.match(r'^\d+([,.]\d+)?$', txt):
|
||
continue
|
||
bbox = line_data.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]]
|
||
cx = sum(xs) / len(xs)
|
||
cy = sum(ys) / len(ys)
|
||
# Определяем границы
|
||
x1, y1, x2, y2 = min(xs), min(ys), max(xs), max(ys)
|
||
numbers.append({
|
||
"text": txt,
|
||
"bbox": bbox,
|
||
"x1": x1, "y1": y1, "x2": x2, "y2": y2,
|
||
"cx": cx, "cy": cy,
|
||
"page": page["page_number"]
|
||
})
|
||
|
||
print(f"[INFO] Всего чисел: {len(numbers)}")
|
||
|
||
# Фильтр 1: исключаем числа из правой части таблиц (x > 0.5w и y > 0.1h)
|
||
# Это эвристика для данного чертежа
|
||
filtered = [n for n in numbers if not (n["x1"] > w * 0.55 and n["y1"] > h * 0.05)]
|
||
print(f"[INFO] После фильтра таблиц: {len(filtered)}")
|
||
|
||
# Фильтр 2: ищем "пары" чисел на одной горизонтали (±15px по Y)
|
||
# Если между числами есть линия — это размерная цепочка
|
||
dimensions = []
|
||
used = set()
|
||
|
||
for i, a in enumerate(filtered):
|
||
if i in used:
|
||
continue
|
||
# Ищем соседей на той же Y
|
||
neighbors = []
|
||
for j, b in enumerate(filtered):
|
||
if i == j or j in used:
|
||
continue
|
||
# Сравниваем Y (горизонтальная линия) или X (вертикальная)
|
||
dy = abs(a["cy"] - b["cy"])
|
||
dx = abs(a["cx"] - b["cx"])
|
||
if dy < 20 and dx > 30 and dx < 600:
|
||
# Проверяем, есть ли между ними тёмная линия
|
||
y_check = int((a["cy"] + b["cy"]) / 2)
|
||
x_start = min(int(a["cx"]), int(b["cx"]))
|
||
x_end = max(int(a["cx"]), int(b["cx"]))
|
||
line_px = img[y_check, x_start:x_end]
|
||
dark_ratio = np.sum(line_px < 200) / len(line_px) if len(line_px) > 0 else 0
|
||
if dark_ratio > 0.3: # >30% тёмных пикселей
|
||
neighbors.append((j, b, dx, "horizontal"))
|
||
|
||
# Ищем вертикальных соседей
|
||
for j, b in enumerate(filtered):
|
||
if i == j or j in used:
|
||
continue
|
||
dx = abs(a["cx"] - b["cx"])
|
||
dy = abs(a["cy"] - b["cy"])
|
||
if dx < 20 and dy > 30 and dy < 600:
|
||
x_check = int((a["cx"] + b["cx"]) / 2)
|
||
y_start = min(int(a["cy"]), int(b["cy"]))
|
||
y_end = max(int(a["cy"]), int(b["cy"]))
|
||
line_px = img[y_start:y_end, x_check]
|
||
dark_ratio = np.sum(line_px < 200) / len(line_px) if len(line_px) > 0 else 0
|
||
if dark_ratio > 0.3:
|
||
neighbors.append((j, b, dy, "vertical"))
|
||
|
||
if neighbors:
|
||
# Берём ближайшего соседа
|
||
neighbors.sort(key=lambda x: x[2])
|
||
j, b, dist, orient = neighbors[0]
|
||
dimensions.append({
|
||
"text": a["text"],
|
||
"bbox": a["bbox"],
|
||
"neighbor_text": b["text"],
|
||
"distance": int(dist),
|
||
"orientation": orient,
|
||
"page": a["page"]
|
||
})
|
||
used.add(i)
|
||
used.add(j)
|
||
|
||
# Одиночные числа — это скорее всего массы/количества из таблиц, игнорируем
|
||
|
||
print(f"[INFO] Размеров найдено: {len(dimensions)}")
|
||
return numbers, dimensions
|
||
|
||
|
||
def visualize(png_path: Path, all_numbers: List[Dict], dimensions: List[Dict], out_path: Path):
|
||
"""Рисует визуализацию: размеры — красные, остальные числа — синие."""
|
||
img = Image.open(png_path)
|
||
draw = ImageDraw.Draw(img)
|
||
|
||
# Все числа (синие)
|
||
dim_texts = {d["text"] for d in dimensions}
|
||
for num in all_numbers:
|
||
bbox = num["bbox"]
|
||
if isinstance(bbox[0], list):
|
||
pts = [(p[0], p[1]) for p in bbox]
|
||
else:
|
||
pts = [(bbox[0], bbox[1]), (bbox[2], bbox[1]), (bbox[2], bbox[3]), (bbox[0], bbox[3])]
|
||
color = "red" if num["text"] in dim_texts else "blue"
|
||
width = 3 if num["text"] in dim_texts else 1
|
||
draw.polygon(pts, outline=color, width=width)
|
||
if num["text"] in dim_texts:
|
||
x = min(p[0] for p in pts)
|
||
y = min(p[1] for p in pts)
|
||
draw.text((x, y-15), num["text"], fill="red")
|
||
|
||
img.save(out_path)
|
||
print(f"[OK] Визуализация сохранена: {out_path}")
|
||
|
||
|
||
def main():
|
||
if len(sys.argv) < 3:
|
||
print("Usage: python dimension_extractor.py <png> <ocr_json>")
|
||
sys.exit(1)
|
||
|
||
png_path = Path(sys.argv[1])
|
||
ocr_path = Path(sys.argv[2])
|
||
out_json = png_path.parent / "dimensions.json"
|
||
out_png = png_path.parent / f"{png_path.stem}_dims_detected.png"
|
||
|
||
all_numbers, dimensions = find_numbers_with_context(ocr_path, png_path)
|
||
|
||
with open(out_json, "w", encoding="utf-8") as f:
|
||
json.dump({
|
||
"dimensions": dimensions,
|
||
"stats": {
|
||
"total_numbers": len(all_numbers),
|
||
"dimensions_found": len(dimensions)
|
||
}
|
||
}, f, ensure_ascii=False, indent=2)
|
||
print(f"[OK] Результаты сохранены: {out_json}")
|
||
|
||
visualize(png_path, all_numbers, dimensions, out_png)
|
||
|
||
print("\nНайденные размеры:")
|
||
for d in dimensions:
|
||
neighbor = f" → {d['neighbor_text']}" if d['neighbor_text'] else ""
|
||
print(f" {d['text']}{neighbor}")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|