opencode/dimension_extractor.py

#!/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()