From b5f7c6327eaffd61005ed2b7e6c6b07763b1c5ea Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=D0=9A=D0=B8=D1=80=D0=B8=D0=BB=D0=BB=20=D0=91=D0=BB=D0=B8?= =?UTF-8?q?=D0=BD=D0=BE=D0=B2?= Date: Mon, 1 Jun 2026 12:29:26 +0300 Subject: [PATCH] Add tiling OCR, preprocess and visualization tools - 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 --- compare_ocr.py | 160 +++++++++++++++++++++++++++ dimension_extractor.py | 193 +++++++++++++++++++++++++++++++++ ocr_qwen.py | 233 ++++++++++++++++++++++++++++++++++++++++ preprocess_for_ocr.py | 51 +++++++++ process_any_pdf.py | 142 +++++++++++++++++++++--- test_qwen_ocr.py | 139 ++++++++++++++++++++++++ tiling_ocr.py | 157 +++++++++++++++++++++++++++ visualize_dimensions.py | 54 ++++++++++ 8 files changed, 1112 insertions(+), 17 deletions(-) create mode 100644 compare_ocr.py create mode 100644 dimension_extractor.py create mode 100644 ocr_qwen.py create mode 100644 preprocess_for_ocr.py create mode 100644 test_qwen_ocr.py create mode 100644 tiling_ocr.py create mode 100644 visualize_dimensions.py diff --git a/compare_ocr.py b/compare_ocr.py new file mode 100644 index 0000000..82cd08d --- /dev/null +++ b/compare_ocr.py @@ -0,0 +1,160 @@ +#!/usr/bin/env python3 +# -*- coding: utf-8 -*- +""" +Визуализация сравнения: обычный OCR vs tiling OCR. +Рисует bbox зелёным (только обычный), красным (только tiling), жёлтым (оба). +""" + +import sys +import json +import re +from pathlib import Path +from PIL import Image, ImageDraw + + +def load_ocr(path: Path): + """Загружает OCR lines из JSON.""" + data = json.loads(path.read_text(encoding="utf-8")) + if "lines" in data: + return data["lines"] + if "pages" in data: + lines = [] + for page in data["pages"]: + lines.extend(page.get("ocr_lines", [])) + return lines + return [] + + +def bbox_center(box): + if isinstance(box[0], list): + xs = [p[0] for p in box] + ys = [p[1] for p in box] + else: + xs = [box[0], box[2]] + ys = [box[1], box[3]] + return sum(xs)/len(xs), sum(ys)/len(ys) + + +def bbox_rect(box): + if isinstance(box[0], list): + xs = [p[0] for p in box] + ys = [p[1] for p in box] + else: + xs = [box[0], box[2]] + ys = [box[1], box[3]] + return min(xs), min(ys), max(xs), max(ys) + + +def find_matches(text: str, list_b, iou_thresh=0.3): + """Находит ближайший совпадающий bbox в list_b по IoU и тексту.""" + matches = [] + for b in list_b: + if b["text"].strip() != text.strip(): + continue + # IoU + ax1, ay1, ax2, ay2 = bbox_rect(a["bbox"] if 'a' in dir() else None) + # ... (упрощённо: сравниваем по центру) + return matches + + +def visualize_comparison(png_path: Path, normal_ocr_path: Path, tiling_ocr_path: Path, out_path: Path): + """Рисует сравнение.""" + img = Image.open(png_path) + draw = ImageDraw.Draw(img) + + normal = load_ocr(normal_ocr_path) + tiling = load_ocr(tiling_ocr_path) + + # Индексы для быстрого поиска + normal_by_text = {} + for n in normal: + txt = n["text"].strip() + if re.match(r'^\d+([,.]\d+)?$', txt): + normal_by_text.setdefault(txt, []).append(n) + + tiling_by_text = {} + for t in tiling: + txt = t["text"].strip() + if re.match(r'^\d+([,.]\d+)?$', txt): + tiling_by_text.setdefault(txt, []).append(t) + + # Классификация + only_normal = [] # зелёный + only_tiling = [] # красный + both = [] # жёлтый + + all_texts = set(normal_by_text.keys()) | set(tiling_by_text.keys()) + + for txt in all_texts: + n_list = normal_by_text.get(txt, []) + t_list = tiling_by_text.get(txt, []) + + # Сопоставляем по минимальному расстоянию центров + used_t = set() + for n in n_list: + cx_n, cy_n = bbox_center(n["bbox"]) + best = None + best_dist = float('inf') + for i, t in enumerate(t_list): + if i in used_t: + continue + cx_t, cy_t = bbox_center(t["bbox"]) + d = ((cx_n - cx_t)**2 + (cy_n - cy_t)**2)**0.5 + if d < best_dist: + best_dist = d + best = i + + if best is not None and best_dist < 100: # совпадение + both.append((n, t_list[best])) + used_t.add(best) + else: + only_normal.append(n) + + for i, t in enumerate(t_list): + if i not in used_t: + only_tiling.append(t) + + # Рисуем + for item in only_normal: + x1, y1, x2, y2 = bbox_rect(item["bbox"]) + draw.rectangle([x1, y1, x2, y2], outline="green", width=3) + + for item in only_tiling: + x1, y1, x2, y2 = bbox_rect(item["bbox"]) + draw.rectangle([x1, y1, x2, y2], outline="red", width=3) + cx, cy = bbox_center(item["bbox"]) + draw.text((cx, cy-15), item["text"], fill="red") + + for n, t in both: + # Используем bbox из tiling (крупнее) + x1, y1, x2, y2 = bbox_rect(t["bbox"]) + draw.rectangle([x1, y1, x2, y2], outline="yellow", width=2) + + img.save(out_path) + print(f"[OK] Сохранено: {out_path}") + print(f" Только обычный (зелёный): {len(only_normal)}") + print(f" Только tiling (красный): {len(only_tiling)}") + print(f" Оба (жёлтый): {len(both)}") + + # Вывод новых чисел + print(f"\nНовые числа от tiling OCR:") + for item in sorted(only_tiling, key=lambda x: x["bbox"][0][1]): + cx, cy = bbox_center(item["bbox"]) + print(f" {item['text']:>10} x={cx:>8.0f} y={cy:>8.0f}") + + +def main(): + if len(sys.argv) < 4: + print("Usage: python compare_ocr.py ") + sys.exit(1) + + png = Path(sys.argv[1]) + normal = Path(sys.argv[2]) + tiling = Path(sys.argv[3]) + out = png.parent / f"{png.stem}_ocr_compare.png" + + visualize_comparison(png, normal, tiling, out) + + +if __name__ == "__main__": + main() diff --git a/dimension_extractor.py b/dimension_extractor.py new file mode 100644 index 0000000..af147be --- /dev/null +++ b/dimension_extractor.py @@ -0,0 +1,193 @@ +#!/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 ") + 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() diff --git a/ocr_qwen.py b/ocr_qwen.py new file mode 100644 index 0000000..c5c0462 --- /dev/null +++ b/ocr_qwen.py @@ -0,0 +1,233 @@ +#!/usr/bin/env python3 +# -*- coding: utf-8 -*- +""" +OCR через Alibaba Cloud qwen-vl-ocr API. + +Использование: + from ocr_qwen import run_ocr + results = run_ocr(image_path) + +Требует DASHSCOPE_API_KEY в .env +""" + +import os +import json +import base64 +import io +from pathlib import Path +from typing import List, Dict, Tuple +from PIL import Image +from openai import OpenAI + +# Загрузить ключ +_API_KEY = None +_BASE_URL = "https://dashscope-intl.aliyuncs.com/compatible-mode/v1" +_MODEL = "qwen-vl-ocr" + +def _load_key(): + global _API_KEY + if _API_KEY: + return _API_KEY + + # Попробовать .env + env_candidates = [ + Path(__file__).parent / ".env", + Path(__file__).parent.parent / ".env", + Path(__file__).parent.parent.parent / ".env", + ] + for env_path in env_candidates: + if env_path.exists(): + for line in env_path.read_text().splitlines(): + if line.startswith("DASHSCOPE_API_KEY="): + _API_KEY = line.split("=", 1)[1].strip() + os.environ["DASHSCOPE_API_KEY"] = _API_KEY + return _API_KEY + + _API_KEY = os.environ.get("DASHSCOPE_API_KEY") + return _API_KEY + + +def resize_image(image_path: Path, max_size: int = 2048) -> Tuple[str, float, Tuple[int, int]]: + """ + Уменьшает изображение до max_size по длинной стороне. + Возвращает: (base64_string, scale_factor, (orig_w, orig_h)) + """ + img = Image.open(image_path) + orig_w, orig_h = img.size + + # Если уже меньше — не менять + if max(orig_w, orig_h) <= max_size: + with open(image_path, "rb") as f: + b64 = base64.b64encode(f.read()).decode("utf-8") + return b64, 1.0, (orig_w, orig_h) + + # Вычислить новый размер + scale = max_size / max(orig_w, orig_h) + new_w = int(orig_w * scale) + new_h = int(orig_h * scale) + + img_resized = img.resize((new_w, new_h), Image.LANCZOS) + + # Сохранить в буфер + buf = io.BytesIO() + img_resized.save(buf, format="PNG") + b64 = base64.b64encode(buf.getvalue()).decode("utf-8") + + return b64, scale, (orig_w, orig_h) + + +def encode_image(image_path: Path) -> str: + with open(image_path, "rb") as f: + return base64.b64encode(f.read()).decode("utf-8") + + +def parse_qwen_response(raw_text: str) -> List[Dict]: + """Парсит JSON из ответа qwen-vl-ocr.""" + import re + text = raw_text.strip() + + # Удалить markdown code blocks ```json ... ``` + if text.startswith("```"): + lines = text.splitlines() + start = 0 + end = len(lines) + for i, line in enumerate(lines): + if line.strip().startswith("```") and start == 0: + start = i + 1 + elif line.strip() == "```" and start > 0: + end = i + break + text = "\n".join(lines[start:end]).strip() + + # Робастный парсинг: извлекаем каждый объект отдельно через regex + results = [] + # Шаблон: {"text": "...", "rotate_rect": [num, num, num, num, num]} + pattern = r'\{\s*"text":\s*"([^"]*)"\s*,\s*"rotate_rect":\s*\[\s*(-?\d+)\s*,\s*(-?\d+)\s*,\s*(-?\d+)\s*,\s*(-?\d+)\s*,\s*(-?\d+)\s*\]\s*\}' + + for match in re.finditer(pattern, text): + txt = match.group(1) + x, y, w, h, angle = int(match.group(2)), int(match.group(3)), int(match.group(4)), int(match.group(5)), int(match.group(6)) + results.append({ + "text": txt, + "rotate_rect": [x, y, w, h, angle] + }) + + if not results: + # Fallback: попробовать стандартный JSON парсинг + try: + json_match = re.search(r'\[[\s\S]*\]', text) + if json_match: + data = json.loads(json_match.group(0)) + if isinstance(data, list): + return data + except Exception: + pass + print(f"[WARN] Regex parser не нашёл объекты, JSON тоже не распарсился") + print(f"[WARN] Text preview: {text[:200]}") + + return results + + +def run_ocr(image_path: Path, verbose: bool = False) -> List[Dict]: + """ + Запускает qwen-vl-ocr на изображении. + + Returns: + Список словарей: { + "text": str, + "bbox": [x1, y1, x2, y2, angle], # rotate_rect format + "confidence": float # estimated + } + """ + api_key = _load_key() + if not api_key: + raise RuntimeError("DASHSCOPE_API_KEY not found in .env or environment") + + client = OpenAI(api_key=api_key, base_url=_BASE_URL) + + # Уменьшить изображение для экономии токенов + b64, scale, (orig_w, orig_h) = resize_image(image_path, max_size=2048) + data_url = f"data:image/png;base64,{b64}" + + if verbose: + orig_size = image_path.stat().st_size / 1024 + print(f"[qwen-ocr] Отправка {image_path.name} (orig {orig_w}x{orig_h}, scale={scale:.2f}, {orig_size:.0f} KB)...", flush=True) + + response = client.chat.completions.create( + model=_MODEL, + messages=[ + { + "role": "user", + "content": [ + { + "type": "text", + "text": ( + "Распознай все текстовые элементы на этом чертеже. " + "Для каждого текста верни ОТДЕЛЬНЫЙ JSON-объект с полями: text, rotate_rect [x,y,w,h,angle]. " + "ВАЖНО: каждый текст — отдельный объект, без дублирующихся ключей в одном объекте. " + "Пример правильного формата:\n" + '[{"text": "Бетон", "rotate_rect": [100, 50, 30, 10, 0]}, {"text": "В30", "rotate_rect": [100, 65, 20, 10, 0]}]' + "\nОтветь строго в формате JSON-массива без markdown." + ), + }, + {"type": "image_url", "image_url": {"url": data_url}}, + ], + } + ], + temperature=0.1, + max_tokens=8192, + ) + + raw = response.choices[0].message.content.strip() + + # Сохранить raw для отладки + debug_path = image_path.parent / f"{image_path.stem}_qwen_raw.txt" + debug_path.write_text(raw, encoding="utf-8") + + items = parse_qwen_response(raw) + + # Конвертировать rotate_rect в наш формат, масштабируя обратно к оригиналу + results = [] + for item in items: + rect = item.get("rotate_rect", [0, 0, 0, 0, 0]) + if len(rect) >= 4: + x, y, w, h = rect[0], rect[1], rect[2], rect[3] + # Масштабировать обратно к оригинальному размеру + if scale != 1.0: + x = round(x / scale) + y = round(y / scale) + w = round(w / scale) + h = round(h / scale) + # bbox: [[x1,y1],[x2,y2],[x3,y3],[x4,y4]] + bbox = [[x, y], [x + w, y], [x + w, y + h], [x, y + h]] + else: + bbox = None + + results.append({ + "text": item.get("text", ""), + "bbox": bbox, + "confidence": 0.95, # qwen-vl-ocr не возвращает confidence, ставим высокий + "source": "qwen-vl-ocr" + }) + + if verbose: + print(f"[qwen-ocr] Найдено {len(results)} элементов") + + return results + + +if __name__ == "__main__": + import sys + if len(sys.argv) < 2: + print("Usage: python ocr_qwen.py ") + sys.exit(1) + + image_path = Path(sys.argv[1]) + results = run_ocr(image_path, verbose=True) + + print(f"\nНайдено {len(results)} текстовых элементов:") + for r in results[:20]: + print(f" '{r['text']}' bbox={r['bbox']}") + + if len(results) > 20: + print(f" ... и ещё {len(results) - 20}") diff --git a/preprocess_for_ocr.py b/preprocess_for_ocr.py new file mode 100644 index 0000000..31255c6 --- /dev/null +++ b/preprocess_for_ocr.py @@ -0,0 +1,51 @@ +#!/usr/bin/env python3 +# -*- coding: utf-8 -*- +""" +Предобработка PNG для улучшения OCR размерных чисел. + +Алгоритм: +1. CLAHE — локальное повышение контраста +2. Unsharp mask — повышение резкости +3. Инвертирование (опционально для некоторых OCR) +4. Масштабирование x2 (если исходное маленькое) +""" + +import sys +from pathlib import Path +import cv2 +import numpy as np + + +def preprocess_for_ocr(img_path: Path, out_path: Path, scale: float = 2.0): + img = cv2.imread(str(img_path), cv2.IMREAD_GRAYSCALE) + if img is None: + raise RuntimeError(f"Cannot load {img_path}") + + # Масштабирование + if scale != 1.0: + h, w = img.shape + img = cv2.resize(img, (int(w*scale), int(h*scale)), interpolation=cv2.INTER_CUBIC) + + # CLAHE (локальный контраст) + clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8,8)) + img = clahe.apply(img) + + # Unsharp mask + gaussian = cv2.GaussianBlur(img, (0,0), 3) + img = cv2.addWeighted(img, 1.5, gaussian, -0.5, 0) + + # Нормализация + img = cv2.normalize(img, None, 0, 255, cv2.NORM_MINMAX) + + cv2.imwrite(str(out_path), img) + print(f"[OK] Предобработка сохранена: {out_path}") + + +if __name__ == "__main__": + if len(sys.argv) < 2: + print("Usage: python preprocess_for_ocr.py ") + sys.exit(1) + + png = Path(sys.argv[1]) + out = png.parent / f"{png.stem}_preproc.png" + preprocess_for_ocr(png, out) diff --git a/process_any_pdf.py b/process_any_pdf.py index 7e6d9c3..b4854a5 100644 --- a/process_any_pdf.py +++ b/process_any_pdf.py @@ -2,14 +2,21 @@ # -*- coding: utf-8 -*- """ Универсальное распознавание PDF в указанную папку. +Поддерживает: + - RapidOCR (локально, быстро) + - RapidOCR + tiling (для больших чертежей) + - qwen-vl-ocr (API, точнее) + Использование: - python process_any_pdf.py + python process_any_pdf.py [--use-qwen] [--use-tiling] """ import sys import json +import re import fitz from pathlib import Path +from PIL import Image from rapidocr_onnxruntime import RapidOCR # ------------------------------------------------------------------ @@ -17,19 +24,95 @@ from rapidocr_onnxruntime import RapidOCR # ------------------------------------------------------------------ DPI = 300 BATCH_SIZE = 5 +TILE_SIZE = 2000 +TILE_OVERLAP = 200 engine = RapidOCR() +# qwen-vl-ocr lazy import +try: + from ocr_qwen import run_ocr as qwen_ocr + QWEN_AVAILABLE = True +except ImportError: + QWEN_AVAILABLE = False + + # ------------------------------------------------------------------ -def process_pdf(pdf_path: Path, out_dir: Path): +# Tiling OCR helpers +# ------------------------------------------------------------------ +def _make_tiles(img: Image.Image, tile_size: int = 2000, overlap: int = 200): + w, h = img.size + tiles = [] + step = tile_size - overlap + for y in range(0, h, step): + for x in range(0, w, step): + x2 = min(x + tile_size, w) + y2 = min(y + tile_size, h) + tiles.append((x, y, img.crop((x, y, x2, y2)))) + return tiles + + +def _bbox_iou(a, b): + def _rect(box): + if isinstance(box[0], list): + xs = [p[0] for p in box] + ys = [p[1] for p in box] + return min(xs), min(ys), max(xs), max(ys) + return box[0], box[1], box[2], box[3] + + ax1, ay1, ax2, ay2 = _rect(a) + bx1, by1, bx2, by2 = _rect(b) + ix1, iy1 = max(ax1, bx1), max(ay1, by1) + ix2, iy2 = min(ax2, bx2), min(ay2, by2) + if ix2 <= ix1 or iy2 <= iy1: + return 0.0 + inter = (ix2 - ix1) * (iy2 - iy1) + union = (ax2 - ax1) * (ay2 - ay1) + (bx2 - bx1) * (by2 - by1) - inter + return inter / union if union > 0 else 0.0 + + +def run_tiling_ocr(img_path: Path, conf_threshold: float = 0.5): + """Запускает RapidOCR по кропам и объединяет результаты.""" + img = Image.open(img_path) + tiles = _make_tiles(img, TILE_SIZE, TILE_OVERLAP) + all_results = [] + for off_x, off_y, crop in tiles: + tmp = f"/tmp/tile_ocr.png" + crop.save(tmp) + res = engine(tmp) + if res and res[0]: + for item in res[0]: + box, txt, score = item + if score < conf_threshold: + continue + shifted = [[pt[0] + off_x, pt[1] + off_y] for pt in box] + all_results.append({"text": txt, "confidence": float(score), "bbox": shifted}) + + # Дедупликация по IoU + unique = [] + for r in sorted(all_results, key=lambda x: -x["confidence"]): + is_dup = any(_bbox_iou(r["bbox"], u["bbox"]) > 0.5 for u in unique) + if not is_dup: + unique.append(r) + return unique + +# ------------------------------------------------------------------ +def process_pdf(pdf_path: Path, out_dir: Path, use_qwen: bool = False, use_tiling: bool = False): out_dir.mkdir(parents=True, exist_ok=True) doc = fitz.open(pdf_path) total = len(doc) - print(f"=== PDF: {pdf_path.name} | Страниц: {total} -> {out_dir} ===\n") + print(f"=== PDF: {pdf_path.name} | Страниц: {total} -> {out_dir} ===") + if use_qwen: + print(f"[INFO] OCR engine: qwen-vl-ocr (API)") + elif use_tiling: + print(f"[INFO] OCR engine: RapidOCR + tiling ({TILE_SIZE}px tiles)") + else: + print(f"[INFO] OCR engine: RapidOCR (local)") + print() all_pages = [] for i in range(total): - print(f"[{i+1}/{total}] Рендер + OCR ...", end=" ") + print(f"[{i+1}/{total}] Рендер + OCR ...", end=" ", flush=True) page = doc.load_page(i) raw_text = page.get_text("text").strip() @@ -38,16 +121,21 @@ def process_pdf(pdf_path: Path, out_dir: Path): img_path = out_dir / f"page_{i+1:03d}.png" pix.save(img_path) - res = engine(img_path) - ocr_lines = [] - if res and res[0] is not None: - for item in res[0]: - box, txt, score = item - ocr_lines.append({ - "text": txt, - "confidence": float(score), - "bbox": box - }) + # Выбор OCR engine + if use_qwen and QWEN_AVAILABLE: + try: + ocr_lines = qwen_ocr(img_path, verbose=False) + print(f"qwen-ocr строк: {len(ocr_lines)}") + except Exception as e: + print(f"qwen-ocr ERR: {e}, fallback to RapidOCR") + ocr_lines = _run_rapidocr(img_path) + print(f"RapidOCR строк: {len(ocr_lines)}") + elif use_tiling: + ocr_lines = run_tiling_ocr(img_path) + print(f"Tiling OCR строк: {len(ocr_lines)}") + else: + ocr_lines = _run_rapidocr(img_path) + print(f"RapidOCR строк: {len(ocr_lines)}") all_pages.append({ "page_number": i + 1, @@ -56,18 +144,38 @@ def process_pdf(pdf_path: Path, out_dir: Path): "ocr_lines": ocr_lines, "ocr_line_count": len(ocr_lines) }) - print(f"OCR строк: {len(ocr_lines)}") if (i + 1) % BATCH_SIZE == 0 or i == total - 1: with open(out_dir / "full_ocr_results.json", "w", encoding="utf-8") as f: json.dump({"pages": all_pages}, f, ensure_ascii=False, indent=2) - print(f" -> промежуточное сохранение ({i+1} страниц)") + print(f" -> сохранено ({i+1} страниц)") doc.close() print(f"\n=== Готово. Результат в {out_dir} ===") + +def _run_rapidocr(img_path: Path): + res = engine(img_path) + ocr_lines = [] + if res and res[0] is not None: + for item in res[0]: + box, txt, score = item + ocr_lines.append({ + "text": txt, + "confidence": float(score), + "bbox": box + }) + return ocr_lines + # ------------------------------------------------------------------ def main(): + use_qwen = "--use-qwen" in sys.argv + use_tiling = "--use-tiling" in sys.argv + if use_qwen: + sys.argv.remove("--use-qwen") + if use_tiling: + sys.argv.remove("--use-tiling") + if len(sys.argv) < 2: pdf_file = "123.pdf" out_name = "output_123" @@ -82,7 +190,7 @@ def main(): print(f"[ERR] Файл не найден: {pdf_path}") sys.exit(1) - process_pdf(pdf_path, out_dir) + process_pdf(pdf_path, out_dir, use_qwen=use_qwen, use_tiling=use_tiling) if __name__ == "__main__": main() diff --git a/test_qwen_ocr.py b/test_qwen_ocr.py new file mode 100644 index 0000000..7e220a7 --- /dev/null +++ b/test_qwen_ocr.py @@ -0,0 +1,139 @@ +#!/usr/bin/env python3 +# -*- coding: utf-8 -*- +""" +Тест Alibaba Cloud DashScope qwen-vl-ocr на чертеже. + +Использование: + python test_qwen_ocr.py +""" + +import os +import sys +import base64 +import json +from pathlib import Path +from openai import OpenAI + +# Загрузить ключ из .env (рядом со скриптом) +env_path = Path(__file__).parent / ".env" +API_KEY = None +if env_path.exists(): + for line in env_path.read_text().splitlines(): + if line.startswith("DASHSCOPE_API_KEY="): + API_KEY = line.split("=", 1)[1].strip() + os.environ["DASHSCOPE_API_KEY"] = API_KEY + break + +if not API_KEY: + API_KEY = os.environ.get("DASHSCOPE_API_KEY") +BASE_URL = "https://dashscope-intl.aliyuncs.com/compatible-mode/v1" +MODEL = "qwen-vl-ocr" + + +def encode_image(image_path: Path) -> str: + with open(image_path, "rb") as f: + return base64.b64encode(f.read()).decode("utf-8") + + +def test_ocr(image_path: Path): + client = OpenAI(api_key=API_KEY, base_url=BASE_URL) + + b64 = encode_image(image_path) + data_url = f"data:image/png;base64,{b64}" + + print(f"Отправляем {image_path.name} в qwen-vl-ocr...") + print(f"Размер файла: {image_path.stat().st_size / 1024 / 1024:.1f} MB") + + response = client.chat.completions.create( + model=MODEL, + messages=[ + { + "role": "user", + "content": [ + { + "type": "text", + "text": ( + "Распознай все текстовые элементы на этом чертеже. " + "Для каждого текста укажи:\n" + "- сам текст\n" + "- координаты bbox (x1,y1,x2,y2)\n" + "- confidence (если доступен)\n" + "Ответь в формате JSON-массива." + ), + }, + {"type": "image_url", "image_url": {"url": data_url}}, + ], + } + ], + temperature=0.1, + max_tokens=2048, + ) + + raw = response.choices[0].message.content + print("\n=== ОТВЕТ МОДЕЛИ ===") + print(raw[:2000]) + print("=" * 50) + + # Сохранить результат + out_path = image_path.parent / f"qwen_ocr_result_{image_path.stem}.json" + with open(out_path, "w", encoding="utf-8") as f: + f.write(raw) + print(f"\n[OK] Сохранено: {out_path}") + + +def describe_image(image_path: Path): + """Просто описание того, что модель видит на чертеже.""" + client = OpenAI(api_key=API_KEY, base_url=BASE_URL) + + b64 = encode_image(image_path) + data_url = f"data:image/png;base64,{b64}" + + print(f"\nОтправляем {image_path.name} на описание...") + + response = client.chat.completions.create( + model=MODEL, + messages=[ + { + "role": "user", + "content": [ + { + "type": "text", + "text": ( + "Опиши подробно, что ты видишь на этом изображении. " + "Чертеж здания или что-то другое? Какие элементы видны? " + "Размеры, текст, линии, оси — всё, что различимо." + ), + }, + {"type": "image_url", "image_url": {"url": data_url}}, + ], + } + ], + temperature=0.3, + max_tokens=1024, + ) + + desc = response.choices[0].message.content + print("\n=== ОПИСАНИЕ ===") + print(desc) + print("=" * 50) + return desc + + +def main(): + if len(sys.argv) < 2: + print("Usage: python test_qwen_ocr.py [--describe]") + sys.exit(1) + + image_path = Path(sys.argv[1]) + if not image_path.exists(): + print(f"[ERR] Файл не найден: {image_path}") + sys.exit(1) + + if "--describe" in sys.argv: + describe_image(image_path) + else: + test_ocr(image_path) + + +if __name__ == "__main__": + main() diff --git a/tiling_ocr.py b/tiling_ocr.py new file mode 100644 index 0000000..9e425f5 --- /dev/null +++ b/tiling_ocr.py @@ -0,0 +1,157 @@ +#!/usr/bin/env python3 +# -*- coding: utf-8 -*- +""" +Tiling OCR для больших чертежей. + +Разрезает PNG на перекрывающиеся кропы, прогоняет OCR на каждом, +объединяет результаты с дедупликацией. + +Эффект: каждый кроп масштабирован "крупнее" для OCR — мелкий текст +находится на бОльшем % площади кропа. +""" + +import sys +import json +import re +from pathlib import Path +from typing import List, Dict, Tuple +from PIL import Image +from rapidocr_onnxruntime import RapidOCR + + +def make_tiles(img: Image.Image, tile_size: int = 2000, overlap: int = 200) -> List[Tuple[int, int, Image.Image]]: + """ + Генерирует кропы с перекрытием. + Возвращает: [(offset_x, offset_y, cropped_image), ...] + """ + w, h = img.size + tiles = [] + step = tile_size - overlap + + for y in range(0, h, step): + for x in range(0, w, step): + x2 = min(x + tile_size, w) + y2 = min(y + tile_size, h) + crop = img.crop((x, y, x2, y2)) + tiles.append((x, y, crop)) + + return tiles + + +def iou_bbox(a: List, b: List) -> float: + """IoU двух bbox в формате [[x1,y1],[x2,y2],[x3,y3],[x4,y4]].""" + def _get_rect(box): + if isinstance(box[0], list): + xs = [p[0] for p in box] + ys = [p[1] for p in box] + return min(xs), min(ys), max(xs), max(ys) + else: + return box[0], box[1], box[2], box[3] + + ax1, ay1, ax2, ay2 = _get_rect(a) + bx1, by1, bx2, by2 = _get_rect(b) + + ix1 = max(ax1, bx1) + iy1 = max(ay1, by1) + ix2 = min(ax2, bx2) + iy2 = min(ay2, by2) + + if ix2 <= ix1 or iy2 <= iy1: + return 0.0 + + inter = (ix2 - ix1) * (iy2 - iy1) + area_a = (ax2 - ax1) * (ay2 - ay1) + area_b = (bx2 - bx1) * (by2 - by1) + union = area_a + area_b - inter + return inter / union if union > 0 else 0.0 + + +def run_tiling_ocr(png_path: Path, tile_size: int = 2000, overlap: int = 200, conf_threshold: float = 0.5): + """Основная функция.""" + print(f"[INFO] Загрузка {png_path.name}...") + img = Image.open(png_path) + print(f"[INFO] Размер: {img.size}") + + tiles = make_tiles(img, tile_size, overlap) + print(f"[INFO] Кропов: {len(tiles)}") + + engine = RapidOCR() + all_results = [] + + for i, (off_x, off_y, crop) in enumerate(tiles, 1): + # Временно сохранить кроп + tmp_path = f"/tmp/tile_{i:03d}.png" + crop.save(tmp_path) + + print(f" [{i}/{len(tiles)}] tile @ ({off_x}, {off_y}) size {crop.size} ...", end=" ", flush=True) + res = engine(tmp_path) + + tile_lines = 0 + if res and res[0]: + for item in res[0]: + box, txt, score = item + if score < conf_threshold: + continue + # Сдвинуть bbox на offset кропа + shifted_box = [] + for pt in box: + shifted_box.append([pt[0] + off_x, pt[1] + off_y]) + all_results.append({ + "text": txt, + "confidence": float(score), + "bbox": shifted_box + }) + tile_lines += 1 + print(f"{tile_lines} lines") + + # Дедупликация: если два bbox пересекаются (IoU > 0.5) — оставляем тот, что с higher confidence + print(f"[INFO] Дедупликация {len(all_results)} строк...") + unique = [] + for r in sorted(all_results, key=lambda x: -x["confidence"]): + is_dup = False + for u in unique: + if iou_bbox(r["bbox"], u["bbox"]) > 0.5: + is_dup = True + break + if not is_dup: + unique.append(r) + + print(f"[OK] Уникальных строк: {len(unique)}") + return unique + + +def main(): + if len(sys.argv) < 2: + print("Usage: python tiling_ocr.py [tile_size] [overlap]") + sys.exit(1) + + png_path = Path(sys.argv[1]) + tile_size = int(sys.argv[2]) if len(sys.argv) > 2 else 2000 + overlap = int(sys.argv[3]) if len(sys.argv) > 3 else 200 + + results = run_tiling_ocr(png_path, tile_size, overlap) + + # Сохранить результаты + out_json = png_path.parent / f"{png_path.stem}_tiling_ocr.json" + with open(out_json, "w", encoding="utf-8") as f: + json.dump({ + "source": str(png_path), + "tile_size": tile_size, + "overlap": overlap, + "total_lines": len(results), + "lines": results + }, f, ensure_ascii=False, indent=2) + print(f"[OK] Сохранено: {out_json}") + + # Вывести числа + nums = [r for r in results if re.match(r'^\d+([,.]\d+)?$', r["text"].strip())] + print(f"\nНайдено {len(nums)} чисел:") + for n in sorted(nums, key=lambda x: x["bbox"][0][1]): + bbox = n["bbox"] + cx = sum(p[0] for p in bbox) / len(bbox) + cy = sum(p[1] for p in bbox) / len(bbox) + print(f" {n['text']:>10} x={cx:>8.0f} y={cy:>8.0f} conf={n['confidence']:.2f}") + + +if __name__ == "__main__": + main() diff --git a/visualize_dimensions.py b/visualize_dimensions.py new file mode 100644 index 0000000..079d868 --- /dev/null +++ b/visualize_dimensions.py @@ -0,0 +1,54 @@ +#!/usr/bin/env python3 +# -*- coding: utf-8 -*- +""" +Визуализация найденных размерных чисел на PNG. +Рисует bbox вокруг чисел, извлечённых из OCR. +""" + +import sys +import json +import re +from pathlib import Path +from PIL import Image, ImageDraw + + +def visualize_dimensions(ocr_json_path: Path, png_path: Path, out_path: Path): + """Рисует bbox вокруг чисел на PNG.""" + ocr = json.loads(ocr_json_path.read_text(encoding="utf-8")) + img = Image.open(png_path) + draw = ImageDraw.Draw(img) + + found = 0 + for page in ocr.get("pages", []): + for line in page.get("ocr_lines", []): + txt = line["text"].strip() + if re.match(r'^\d+([,.]\d+)?$', txt): + bbox = line.get("bbox") + if bbox: + # bbox: [[x1,y1],[x2,y2],[x3,y3],[x4,y4]] + 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])] + draw.polygon(pts, outline="red", width=3) + # Подпись + x = min(p[0] for p in pts) + y = min(p[1] for p in pts) + draw.text((x, y-20), txt, fill="red") + found += 1 + + img.save(out_path) + print(f"[OK] Найдено {found} размерных чисел. Сохранено: {out_path}") + + +if __name__ == "__main__": + if len(sys.argv) < 3: + print("Usage: python visualize_dimensions.py ") + sys.exit(1) + + ocr_json = Path(sys.argv[1]) + png = Path(sys.argv[2]) + out = png.parent / f"{png.stem}_dims.png" + + visualize_dimensions(ocr_json, png, out)