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