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
This commit is contained in:
parent
c756a5766b
commit
b5f7c6327e
160
compare_ocr.py
Normal file
160
compare_ocr.py
Normal file
@ -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 <png> <normal_ocr.json> <tiling_ocr.json>")
|
||||
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()
|
||||
193
dimension_extractor.py
Normal file
193
dimension_extractor.py
Normal file
@ -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 <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()
|
||||
233
ocr_qwen.py
Normal file
233
ocr_qwen.py
Normal file
@ -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 <image.png>")
|
||||
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}")
|
||||
51
preprocess_for_ocr.py
Normal file
51
preprocess_for_ocr.py
Normal file
@ -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 <png>")
|
||||
sys.exit(1)
|
||||
|
||||
png = Path(sys.argv[1])
|
||||
out = png.parent / f"{png.stem}_preproc.png"
|
||||
preprocess_for_ocr(png, out)
|
||||
@ -2,14 +2,21 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""
|
||||
Универсальное распознавание PDF в указанную папку.
|
||||
Поддерживает:
|
||||
- RapidOCR (локально, быстро)
|
||||
- RapidOCR + tiling (для больших чертежей)
|
||||
- qwen-vl-ocr (API, точнее)
|
||||
|
||||
Использование:
|
||||
python process_any_pdf.py <pdf_file> <output_folder_name>
|
||||
python process_any_pdf.py <pdf_file> <output_folder> [--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()
|
||||
|
||||
139
test_qwen_ocr.py
Normal file
139
test_qwen_ocr.py
Normal file
@ -0,0 +1,139 @@
|
||||
#!/usr/bin/env python3
|
||||
# -*- coding: utf-8 -*-
|
||||
"""
|
||||
Тест Alibaba Cloud DashScope qwen-vl-ocr на чертеже.
|
||||
|
||||
Использование:
|
||||
python test_qwen_ocr.py <png_file>
|
||||
"""
|
||||
|
||||
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 <png_file> [--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()
|
||||
157
tiling_ocr.py
Normal file
157
tiling_ocr.py
Normal file
@ -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 <png> [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()
|
||||
54
visualize_dimensions.py
Normal file
54
visualize_dimensions.py
Normal file
@ -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 <ocr_json> <png>")
|
||||
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)
|
||||
Loading…
Reference in New Issue
Block a user