#!/usr/bin/env python3 # -*- coding: utf-8 -*- """ VLM Describer — объективное извлечение структуры чертежа. Отправляет PNG в qwen-vl-plus (DashScope API) с промптом на фактическое описание содержимого. НЕ ищет ошибки, НЕ оценивает качество. Результат: /vlm_extraction.json — структурированное описание каждой страницы для использования в RAG и cross-verification. Использование: python vlm_describer.py [--model MODEL] Требует DASHSCOPE_API_KEY в .env или окружении. """ import os import sys import json import base64 import io import re 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" DEFAULT_MODEL = "qwen-vl-plus" def _load_api_key(): global API_KEY if API_KEY: return API_KEY env_candidates = [ Path(__file__).parent / ".env", Path(__file__).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 EXTRACTION_PROMPT = ( "Ты — система распознавания чертежей. Опиши объективно, что изображено на этой странице. " "НЕ ищи ошибки, НЕ оценивай качество. Просто перечисли факты.\n\n" "Ответь СТРОГО в формате JSON (без markdown):\n" "{\n" ' "page_type": "plan / section / elevation / specification / detail / general_view / table / unknown",\n' ' "title": "заголовок или null",\n' ' "beams": ["Балка Б-1"],\n' ' "positions": ["П-1"],\n' ' "gosts": ["ГОСТ ..."],\n' ' "description": "2-3 предложения о содержимом"\n' "}\n\n" "ПРАВИЛА:\n" "- Только реальные элементы с чертежа, не придумывай\n" "- Пустой массив [] если нет элементов данного типа\n" "- НЕ включай массы из таблиц в размеры\n" "- Описание — только факты, без оценок" ) def resize_image(image_path: Path, max_size: int = 2048) -> Tuple[str, float, Tuple[int, int]]: 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 parse_json_response(text: str) -> Dict: """Парсит JSON из ответа VLM.""" text = text.strip() if text.startswith("```"): text = re.sub(r"^```[a-zA-Z]*\n?", "", text) text = re.sub(r"\n?```$", "", text) text = text.strip() json_match = re.search(r'\{[\s\S]*\}', text) if json_match: text = json_match.group(0) try: return json.loads(text) except json.JSONDecodeError as e: print(f"[WARN] Не удалось распарсить JSON: {e}") print(f"[WARN] Raw preview: {text[:500]}") return { "page_type": "unknown", "title": None, "elements": [], "beams": [], "positions": [], "dimensions": [], "gosts": [], "tables": [], "description": text[:500] if text else "", "parse_error": str(e) } def describe_page(image_path: Path, model: str) -> Dict: """Отправляет PNG в qwen-vl API, получает структурированное описание.""" api_key = _load_api_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}" response = client.chat.completions.create( model=model, messages=[ { "role": "user", "content": [ {"type": "text", "text": EXTRACTION_PROMPT}, {"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}_vlm_raw.txt" debug_path.write_text(raw, encoding="utf-8") result = parse_json_response(raw) result["_meta"] = { "image": image_path.name, "original_size": [orig_w, orig_h], "scale": scale, } return result def run_vlm_describer(folder: Path, model: str = DEFAULT_MODEL): """Запускает VLM Describer для всех PNG в папке.""" png_files = sorted(folder.glob("page_*.png")) if not png_files: print(f"[ERR] В папке {folder} не найдены page_*.png") sys.exit(1) out_path = folder / "vlm_extraction.json" extractions = {} print(f"[INFO] VLM Describer: {len(png_files)} страниц") print(f"[INFO] API: DashScope ({BASE_URL})") print(f"[INFO] Модель: {model}\n") for i, png in enumerate(png_files, 1): print(f"[{i}/{len(png_files)}] {png.name} ...", end=" ", flush=True) try: data = describe_page(png, model) extractions[png.name] = data elem_count = len(data.get("beams", [])) + len(data.get("positions", [])) + len(data.get("gosts", [])) print(f"OK ({elem_count} элементов)") except Exception as e: print(f"ERR: {e}") extractions[png.name] = { "page_type": "unknown", "error": str(e), "elements": [], "beams": [], "positions": [], "dimensions": [], "gosts": [], "tables": [], "description": "" } with open(out_path, "w", encoding="utf-8") as f: json.dump(extractions, f, ensure_ascii=False, indent=2) total_elems = sum( len(v.get("beams", [])) + len(v.get("positions", [])) + len(v.get("gosts", [])) for v in extractions.values() ) print(f"\n[OK] VLM extraction сохранён: {out_path}") print(f" Страниц: {len(png_files)}, Всего элементов: {total_elems}") def main(): import argparse parser = argparse.ArgumentParser(description="VLM Describer для чертежей") parser.add_argument("folder", help="Папка с page_*.png") parser.add_argument("--model", default=DEFAULT_MODEL, help="Имя модели (default: qwen-vl-plus)") args = parser.parse_args() folder = Path(args.folder) run_vlm_describer(folder, args.model) if __name__ == "__main__": main()