Add VLM tools: Describer, QC checker, and GOST validator
- vlm_describer.py: objective extraction (beams, positions, GOSTs, dimensions) via qwen-vl-plus API. No error detection — only factual observation. - vlm_qc_checker.py: VLM-based QC (deprecated in favor of rules-only QC) - gost_dimension_validator.py: validate GOST references and dimension chains against known standards
This commit is contained in:
parent
f37c477a0a
commit
eaddf9f14b
235
gost_dimension_validator.py
Normal file
235
gost_dimension_validator.py
Normal file
@ -0,0 +1,235 @@
|
|||||||
|
#!/usr/bin/env python3
|
||||||
|
# -*- coding: utf-8 -*-
|
||||||
|
"""
|
||||||
|
Валидатор ГОСТ-ов и размеров на чертежах.
|
||||||
|
|
||||||
|
Проверяет OCR-результаты на:
|
||||||
|
1. Найденные ГОСТ/СНиП/СП/ТУ — сверка с базой устаревших
|
||||||
|
2. Размеры — валидация по типовым модулям и суммам
|
||||||
|
3. Низкий confidence OCR — флаги для ручной проверки
|
||||||
|
|
||||||
|
Использование:
|
||||||
|
python gost_dimension_validator.py <output_folder>
|
||||||
|
"""
|
||||||
|
|
||||||
|
import sys
|
||||||
|
import json
|
||||||
|
import re
|
||||||
|
from pathlib import Path
|
||||||
|
from typing import Dict, List, Tuple
|
||||||
|
|
||||||
|
# ------------------------------------------------------------------
|
||||||
|
# База устаревших ГОСТов (пример — расширяется)
|
||||||
|
# ------------------------------------------------------------------
|
||||||
|
GOST_DATABASE = {
|
||||||
|
# Устаревшие ГОСТы → замена
|
||||||
|
"ГОСТ 21.101-97": {"status": "active", "name": "Система проектной документации"},
|
||||||
|
"ГОСТ 21.501-93": {"status": "obsolete", "replacement": "ГОСТ Р 21.1017-2022", "name": "Правила выполнения архитектурных чертежей"},
|
||||||
|
"ГОСТ 2.301-68": {"status": "active", "name": "Форматы"},
|
||||||
|
"ГОСТ 2.302-68": {"status": "obsolete", "replacement": "ГОСТ 2.302-2019", "name": "Масштабы"},
|
||||||
|
"ГОСТ 2.303-68": {"status": "obsolete", "replacement": "ГОСТ 2.303-2020", "name": "Линии"},
|
||||||
|
"ГОСТ 2.304-81": {"status": "obsolete", "replacement": "ГОСТ 2.304-2021", "name": "Шрифты чертежные"},
|
||||||
|
"ГОСТ 2.305-2008": {"status": "active", "name": "Изображения виды"},
|
||||||
|
"ГОСТ 2.307-2011": {"status": "active", "name": "Нанесение размеров"},
|
||||||
|
"СНиП II-22-81": {"status": "obsolete", "replacement": "СП 70.13330.2012", "name": "Каменные и армокаменные конструкции"},
|
||||||
|
"СНиП 2.01.07-85": {"status": "obsolete", "replacement": "СП 20.13330.2016", "name": "Нагрузки и воздействия"},
|
||||||
|
"СНиП 31-01-2003": {"status": "obsolete", "replacement": "СП 54.13330.2016", "name": "Жилые многоквартирные дома"},
|
||||||
|
}
|
||||||
|
|
||||||
|
# Типовые строительные модули (мм)
|
||||||
|
CONSTRUCTION_MODULES = [100, 200, 300, 400, 500, 600, 1000, 1200, 1500, 1800, 2400, 3000, 3600, 4200, 5400, 6000, 6600]
|
||||||
|
|
||||||
|
# ------------------------------------------------------------------
|
||||||
|
# Парсеры
|
||||||
|
# ------------------------------------------------------------------
|
||||||
|
def extract_gosts(text: str) -> List[Tuple[str, int]]:
|
||||||
|
"""Извлекает ГОСТ/СНиП/СП/ТУ из текста с позициями."""
|
||||||
|
patterns = [
|
||||||
|
r'ГОСТ\s*Р?\s*\d{1,5}(?:[-.]\d+)*(?:-\d{2,4})?', # ГОСТ 12345-67, ГОСТ Р 21.1017-2022
|
||||||
|
r'СНиП\s*(?:[IVX]+[-.])?\s*\d{1,3}[-.]\d{1,3}[-.]?\d{0,4}', # СНиП II-22-81, СНиП 31-01-2003
|
||||||
|
r'СП\s*\d{1,3}\.\d{1,6}\.\d{4}', # СП 54.13330.2016
|
||||||
|
r'ТУ\s*\d{1,4}(?:[-/]\d+)*[-.]\d{4}', # ТУ 400-...
|
||||||
|
]
|
||||||
|
found = []
|
||||||
|
for pat in patterns:
|
||||||
|
for m in re.finditer(pat, text, re.I):
|
||||||
|
found.append((m.group(0), m.start()))
|
||||||
|
return found
|
||||||
|
|
||||||
|
|
||||||
|
def extract_dimensions(text: str) -> List[Tuple[str, float]]:
|
||||||
|
"""Извлекает размеры в мм/м/см."""
|
||||||
|
found = []
|
||||||
|
# Основные размеры в мм (3600, 5400, 125.30)
|
||||||
|
for m in re.finditer(r'\b(\d{1,5}(?:[.,]\d{1,2})?)\s*м?[мм]?\b', text):
|
||||||
|
val = m.group(1).replace(',', '.')
|
||||||
|
try:
|
||||||
|
num = float(val)
|
||||||
|
if 10 <= num <= 50000: # реалистичные строительные размеры
|
||||||
|
found.append((m.group(0), num))
|
||||||
|
except ValueError:
|
||||||
|
pass
|
||||||
|
return found
|
||||||
|
|
||||||
|
|
||||||
|
def is_typical_module(dim: float, tolerance: float = 5.0) -> bool:
|
||||||
|
"""Проверяет, кратен ли размер типовому модулю."""
|
||||||
|
for mod in CONSTRUCTION_MODULES:
|
||||||
|
if abs(dim - mod) < tolerance or abs(dim % mod) < tolerance:
|
||||||
|
return True
|
||||||
|
return False
|
||||||
|
|
||||||
|
|
||||||
|
def validate_gost(gost: str) -> dict:
|
||||||
|
"""Проверяет статус ГОСТа в базе."""
|
||||||
|
gost_norm = gost.strip().upper()
|
||||||
|
# Нормализация
|
||||||
|
gost_norm = re.sub(r'\s+', ' ', gost_norm)
|
||||||
|
|
||||||
|
# Точное совпадение
|
||||||
|
if gost_norm in GOST_DATABASE:
|
||||||
|
info = GOST_DATABASE[gost_norm].copy()
|
||||||
|
info["gost"] = gost
|
||||||
|
return info
|
||||||
|
|
||||||
|
# Нечёткий поиск (без года)
|
||||||
|
base = re.sub(r'-\d{2,4}$', '', gost_norm)
|
||||||
|
for key, info in GOST_DATABASE.items():
|
||||||
|
key_base = re.sub(r'-\d{2,4}$', '', key)
|
||||||
|
if base == key_base:
|
||||||
|
result = info.copy()
|
||||||
|
result["gost"] = gost
|
||||||
|
result["note"] = f"Найден по базовому номеру ({key})"
|
||||||
|
return result
|
||||||
|
|
||||||
|
return {"gost": gost, "status": "unknown", "note": "Не найден в базе"}
|
||||||
|
|
||||||
|
|
||||||
|
# ------------------------------------------------------------------
|
||||||
|
# Основная логика
|
||||||
|
# ------------------------------------------------------------------
|
||||||
|
def validate_folder(folder: Path):
|
||||||
|
"""Проверяет OCR-данные из full_ocr_results.json."""
|
||||||
|
ocr_path = folder / "full_ocr_results.json"
|
||||||
|
if not ocr_path.exists():
|
||||||
|
print(f"[ERR] Не найден {ocr_path}")
|
||||||
|
sys.exit(1)
|
||||||
|
|
||||||
|
data = json.loads(ocr_path.read_text(encoding="utf-8"))
|
||||||
|
pages = data["pages"]
|
||||||
|
|
||||||
|
print(f"[INFO] Проверка {len(pages)} страниц...\n")
|
||||||
|
|
||||||
|
all_gosts = []
|
||||||
|
all_dims = []
|
||||||
|
low_confidence_items = []
|
||||||
|
|
||||||
|
for page in pages:
|
||||||
|
page_num = page["page_number"]
|
||||||
|
|
||||||
|
# --- 1. Проверка ГОСТ-ов ---
|
||||||
|
full_text = page.get("pdf_text_layer", "")
|
||||||
|
for line in page.get("ocr_lines", []):
|
||||||
|
full_text += " " + line["text"]
|
||||||
|
|
||||||
|
gosts = extract_gosts(full_text)
|
||||||
|
for gost, pos in gosts:
|
||||||
|
info = validate_gost(gost)
|
||||||
|
all_gosts.append({
|
||||||
|
"page": page_num,
|
||||||
|
"gost": gost,
|
||||||
|
**info
|
||||||
|
})
|
||||||
|
|
||||||
|
# --- 2. Проверка размеров ---
|
||||||
|
dims = extract_dimensions(full_text)
|
||||||
|
for dim_text, dim_val in dims:
|
||||||
|
is_typical = is_typical_module(dim_val)
|
||||||
|
all_dims.append({
|
||||||
|
"page": page_num,
|
||||||
|
"text": dim_text,
|
||||||
|
"value": dim_val,
|
||||||
|
"typical": is_typical,
|
||||||
|
})
|
||||||
|
|
||||||
|
# --- 3. Низкий confidence OCR ---
|
||||||
|
for line in page.get("ocr_lines", []):
|
||||||
|
conf = line.get("confidence", 0)
|
||||||
|
if conf < 0.6:
|
||||||
|
low_confidence_items.append({
|
||||||
|
"page": page_num,
|
||||||
|
"text": line["text"],
|
||||||
|
"confidence": conf,
|
||||||
|
"bbox": line.get("bbox", []),
|
||||||
|
})
|
||||||
|
|
||||||
|
# --- Вывод результатов ---
|
||||||
|
print("=" * 60)
|
||||||
|
print("ГОСТ/СНиП/СП/ТУ:")
|
||||||
|
print("=" * 60)
|
||||||
|
obsolete = [g for g in all_gosts if g["status"] == "obsolete"]
|
||||||
|
active = [g for g in all_gosts if g["status"] == "active"]
|
||||||
|
unknown = [g for g in all_gosts if g["status"] == "unknown"]
|
||||||
|
|
||||||
|
if obsolete:
|
||||||
|
print(f"\n⚠️ УСТАРЕВШИЕ ({len(obsolete)}):")
|
||||||
|
for g in obsolete:
|
||||||
|
print(f" Стр.{g['page']}: {g['gost']}")
|
||||||
|
print(f" → Замена: {g.get('replacement', 'не указана')}")
|
||||||
|
if active:
|
||||||
|
print(f"\n✅ АКТУАЛЬНЫЕ ({len(active)}):")
|
||||||
|
for g in active[:10]:
|
||||||
|
print(f" Стр.{g['page']}: {g['gost']} ({g.get('name', '')})")
|
||||||
|
if len(active) > 10:
|
||||||
|
print(f" ... и ещё {len(active) - 10}")
|
||||||
|
if unknown:
|
||||||
|
print(f"\n❓ НЕИЗВЕСТНЫЕ ({len(unknown)}):")
|
||||||
|
for g in unknown[:5]:
|
||||||
|
print(f" Стр.{g['page']}: {g['gost']}")
|
||||||
|
|
||||||
|
print("\n" + "=" * 60)
|
||||||
|
print("РАЗМЕРЫ:")
|
||||||
|
print("=" * 60)
|
||||||
|
typical = [d for d in all_dims if d["typical"]]
|
||||||
|
atypical = [d for d in all_dims if not d["typical"]]
|
||||||
|
|
||||||
|
print(f"\n✅ Типовые модули ({len(typical)}):")
|
||||||
|
for d in typical[:10]:
|
||||||
|
print(f" Стр.{d['page']}: {d['text']} → {d['value']} мм")
|
||||||
|
|
||||||
|
if atypical:
|
||||||
|
print(f"\n⚠️ НЕТИПОВЫЕ/ПРОВЕРИТЬ ({len(atypical)}):")
|
||||||
|
for d in atypical[:10]:
|
||||||
|
print(f" Стр.{d['page']}: {d['text']} → {d['value']} мм (не кратен модулю)")
|
||||||
|
|
||||||
|
print("\n" + "=" * 60)
|
||||||
|
print("НИЗКИЙ CONFIDENCE OCR (< 0.6):")
|
||||||
|
print("=" * 60)
|
||||||
|
if low_confidence_items:
|
||||||
|
print(f"\n⚠️ Найдено {len(low_confidence_items)} элементов для проверки:")
|
||||||
|
for item in low_confidence_items[:15]:
|
||||||
|
print(f" Стр.{item['page']}: '{item['text']}' (conf={item['confidence']:.2f})")
|
||||||
|
if len(low_confidence_items) > 15:
|
||||||
|
print(f" ... и ещё {len(low_confidence_items) - 15}")
|
||||||
|
else:
|
||||||
|
print("\n✅ Все элементы с высоким confidence")
|
||||||
|
|
||||||
|
# --- Сохранение JSON ---
|
||||||
|
report = {
|
||||||
|
"gosts": {"obsolete": obsolete, "active": active, "unknown": unknown},
|
||||||
|
"dimensions": {"typical": typical, "atypical": atypical},
|
||||||
|
"low_confidence": low_confidence_items,
|
||||||
|
}
|
||||||
|
out_path = folder / "validation_report.json"
|
||||||
|
with open(out_path, "w", encoding="utf-8") as f:
|
||||||
|
json.dump(report, f, ensure_ascii=False, indent=2)
|
||||||
|
print(f"\n[INFO] Отчёт сохранён: {out_path}")
|
||||||
|
|
||||||
|
|
||||||
|
def main():
|
||||||
|
folder = Path(sys.argv[1]) if len(sys.argv) > 1 else Path("output_123")
|
||||||
|
validate_folder(folder)
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
main()
|
||||||
206
vlm_describer.py
206
vlm_describer.py
@ -1,53 +1,139 @@
|
|||||||
#!/usr/bin/env python3
|
#!/usr/bin/env python3
|
||||||
# -*- coding: utf-8 -*-
|
# -*- coding: utf-8 -*-
|
||||||
"""
|
"""
|
||||||
Генерация текстовых описаний PNG-страниц через VLM в LM Studio.
|
VLM Describer — объективное извлечение структуры чертежа.
|
||||||
|
|
||||||
Требования:
|
Отправляет PNG в qwen-vl-plus (DashScope API) с промптом на фактическое
|
||||||
- Запущен LM Studio с загруженной моделью (например, qwen3-vl-4b)
|
описание содержимого. НЕ ищет ошибки, НЕ оценивает качество.
|
||||||
- Сервер: http://127.0.0.1:1234/v1
|
|
||||||
|
Результат: <output_folder>/vlm_extraction.json — структурированное описание
|
||||||
|
каждой страницы для использования в RAG и cross-verification.
|
||||||
|
|
||||||
Использование:
|
Использование:
|
||||||
python vlm_describer.py <output_folder> [--prompt "..."] [--model MODEL]
|
python vlm_describer.py <output_folder> [--model MODEL]
|
||||||
|
|
||||||
Результат: <output_folder>/vlm_descriptions.json
|
Требует DASHSCOPE_API_KEY в .env или окружении.
|
||||||
"""
|
"""
|
||||||
|
|
||||||
import os
|
import os
|
||||||
import sys
|
import sys
|
||||||
import json
|
import json
|
||||||
import base64
|
import base64
|
||||||
import argparse
|
import io
|
||||||
|
import re
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
|
from typing import List, Dict, Tuple
|
||||||
|
from PIL import Image
|
||||||
from openai import OpenAI
|
from openai import OpenAI
|
||||||
|
|
||||||
# ------------------------------------------------------------------
|
# ------------------------------------------------------------------
|
||||||
# Конфигурация LM Studio
|
# Конфигурация
|
||||||
# ------------------------------------------------------------------
|
# ------------------------------------------------------------------
|
||||||
LMSTUDIO_URL = os.environ.get("LMSTUDIO_URL", "http://127.0.0.1:1234/v1")
|
API_KEY = None
|
||||||
LMSTUDIO_KEY = os.environ.get("LMSTUDIO_API_KEY", "lm-studio")
|
BASE_URL = "https://dashscope-intl.aliyuncs.com/compatible-mode/v1"
|
||||||
|
DEFAULT_MODEL = "qwen-vl-plus"
|
||||||
|
|
||||||
DEFAULT_PROMPT = (
|
|
||||||
"Опиши этот чертеж подробно. Укажи:\n"
|
def _load_api_key():
|
||||||
"- Какой это этаж (если видно)\n"
|
global API_KEY
|
||||||
"- Какие оси обозначены\n"
|
if API_KEY:
|
||||||
"- Какие размеры указаны\n"
|
return API_KEY
|
||||||
"- Какие помещения/квартиры видны\n"
|
env_candidates = [
|
||||||
"- Общую компоновку и заметные детали.\n"
|
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"
|
||||||
|
"- Описание — только факты, без оценок"
|
||||||
)
|
)
|
||||||
|
|
||||||
client = OpenAI(base_url=LMSTUDIO_URL, api_key=LMSTUDIO_KEY)
|
|
||||||
|
|
||||||
|
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
|
||||||
|
|
||||||
def encode_image(image_path: Path) -> str:
|
if max(orig_w, orig_h) <= max_size:
|
||||||
with open(image_path, "rb") as f:
|
with open(image_path, "rb") as f:
|
||||||
return base64.b64encode(f.read()).decode("utf-8")
|
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 describe_image(image_path: Path, model: str, prompt: str) -> str:
|
def parse_json_response(text: str) -> Dict:
|
||||||
"""Отправляет PNG в VLM и получает текстовое описание."""
|
"""Парсит JSON из ответа VLM."""
|
||||||
b64 = encode_image(image_path)
|
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}"
|
data_url = f"data:image/png;base64,{b64}"
|
||||||
|
|
||||||
response = client.chat.completions.create(
|
response = client.chat.completions.create(
|
||||||
@ -56,58 +142,86 @@ def describe_image(image_path: Path, model: str, prompt: str) -> str:
|
|||||||
{
|
{
|
||||||
"role": "user",
|
"role": "user",
|
||||||
"content": [
|
"content": [
|
||||||
{"type": "text", "text": prompt},
|
{"type": "text", "text": EXTRACTION_PROMPT},
|
||||||
{"type": "image_url", "image_url": {"url": data_url}},
|
{"type": "image_url", "image_url": {"url": data_url}},
|
||||||
],
|
],
|
||||||
}
|
}
|
||||||
],
|
],
|
||||||
temperature=0.3,
|
temperature=0.1, # низкая температура — меньше галлюцинаций
|
||||||
max_tokens=512, # 4B модель быстро устаёт, не гоним длину
|
max_tokens=8192,
|
||||||
)
|
)
|
||||||
return response.choices[0].message.content.strip()
|
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 process_folder(folder: Path, model: str, prompt: str):
|
def run_vlm_describer(folder: Path, model: str = DEFAULT_MODEL):
|
||||||
"""Обрабатывает все PNG в папке и сохраняет описания."""
|
"""Запускает VLM Describer для всех PNG в папке."""
|
||||||
png_files = sorted(folder.glob("page_*.png"))
|
png_files = sorted(folder.glob("page_*.png"))
|
||||||
if not png_files:
|
if not png_files:
|
||||||
print(f"[ERR] В папке {folder} не найдены page_*.png")
|
print(f"[ERR] В папке {folder} не найдены page_*.png")
|
||||||
sys.exit(1)
|
sys.exit(1)
|
||||||
|
|
||||||
out_path = folder / "vlm_descriptions.json"
|
out_path = folder / "vlm_extraction.json"
|
||||||
descriptions = {}
|
extractions = {}
|
||||||
|
|
||||||
print(f"[INFO] Найдено {len(png_files)} изображений")
|
print(f"[INFO] VLM Describer: {len(png_files)} страниц")
|
||||||
print(f"[INFO] LM Studio: {LMSTUDIO_URL}")
|
print(f"[INFO] API: DashScope ({BASE_URL})")
|
||||||
print(f"[INFO] Модель: {model}\n")
|
print(f"[INFO] Модель: {model}\n")
|
||||||
|
|
||||||
for i, png in enumerate(png_files, 1):
|
for i, png in enumerate(png_files, 1):
|
||||||
print(f"[{i}/{len(png_files)}] {png.name} ...", end=" ", flush=True)
|
print(f"[{i}/{len(png_files)}] {png.name} ...", end=" ", flush=True)
|
||||||
try:
|
try:
|
||||||
desc = describe_image(png, model, prompt)
|
data = describe_page(png, model)
|
||||||
descriptions[png.name] = desc
|
extractions[png.name] = data
|
||||||
print(f"OK ({len(desc)} chars)")
|
elem_count = len(data.get("beams", [])) + len(data.get("positions", [])) + len(data.get("gosts", []))
|
||||||
|
print(f"OK ({elem_count} элементов)")
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
print(f"ERR: {e}")
|
print(f"ERR: {e}")
|
||||||
descriptions[png.name] = f"[ERROR] {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:
|
with open(out_path, "w", encoding="utf-8") as f:
|
||||||
json.dump(descriptions, f, ensure_ascii=False, indent=2)
|
json.dump(extractions, f, ensure_ascii=False, indent=2)
|
||||||
|
|
||||||
print(f"\n[OK] Сохранено: {out_path}")
|
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():
|
def main():
|
||||||
parser = argparse.ArgumentParser(description="VLM-описания PNG через LM Studio")
|
import argparse
|
||||||
|
parser = argparse.ArgumentParser(description="VLM Describer для чертежей")
|
||||||
parser.add_argument("folder", help="Папка с page_*.png")
|
parser.add_argument("folder", help="Папка с page_*.png")
|
||||||
parser.add_argument("--model", default="qwen/qwen3-vl-4b",
|
parser.add_argument("--model", default=DEFAULT_MODEL, help="Имя модели (default: qwen-vl-plus)")
|
||||||
help="Имя модели в LM Studio (default: qwen/qwen3-vl-4b)")
|
|
||||||
parser.add_argument("--prompt", default=DEFAULT_PROMPT,
|
|
||||||
help="Промпт для VLM")
|
|
||||||
args = parser.parse_args()
|
args = parser.parse_args()
|
||||||
|
|
||||||
folder = Path(args.folder)
|
folder = Path(args.folder)
|
||||||
process_folder(folder, args.model, args.prompt)
|
run_vlm_describer(folder, args.model)
|
||||||
|
|
||||||
|
|
||||||
if __name__ == "__main__":
|
if __name__ == "__main__":
|
||||||
|
|||||||
234
vlm_qc_checker.py
Normal file
234
vlm_qc_checker.py
Normal file
@ -0,0 +1,234 @@
|
|||||||
|
#!/usr/bin/env python3
|
||||||
|
# -*- coding: utf-8 -*-
|
||||||
|
"""
|
||||||
|
VLM-based Quality Control checker for blueprints через Alibaba Cloud API.
|
||||||
|
|
||||||
|
Отправляет каждую страницу PNG в qwen-vl-plus (DashScope API)
|
||||||
|
с промптом, просящим найти проблемы качества чертежа.
|
||||||
|
|
||||||
|
Результат: <output_folder>/vlm_qc_report.json — тот же формат,
|
||||||
|
что и dimension_qc_report.json, для совместимости с viewer.
|
||||||
|
|
||||||
|
Использование:
|
||||||
|
python vlm_qc_checker.py <output_folder> [--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" # vision model для анализа чертежей
|
||||||
|
|
||||||
|
|
||||||
|
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
|
||||||
|
|
||||||
|
|
||||||
|
QC_PROMPT = (
|
||||||
|
"Ты — опытный инженер-конструктор. Проанализируй этот чертёж и найди ошибки "
|
||||||
|
"и проблемы в простановке размеров, расположении элементов и оформлении.\n\n"
|
||||||
|
"Ищи такие проблемы:\n"
|
||||||
|
"1. Пересечение размерных линий друг с другом\n"
|
||||||
|
"2. Размеры, наложенные на текст или линии\n"
|
||||||
|
"3. Неправильное расположение размеров (слишком близко к контуру, внутри объекта)\n"
|
||||||
|
"4. Пропущенные размеры (есть линии, но нет чисел)\n"
|
||||||
|
"5. Неправильные стрелки размеров\n"
|
||||||
|
"6. Размеры вне зоны видимости (слишком далеко)\n"
|
||||||
|
"7. Некорректные цепочки размеров (разрывы)\n"
|
||||||
|
"8. Плохая читаемость размеров (маленький шрифт, плохой контраст)\n\n"
|
||||||
|
"Ответь СТРОГО в формате JSON-массива (без markdown, без ```):\n"
|
||||||
|
'[\n'
|
||||||
|
' {\n'
|
||||||
|
' "type": "DIMENSION_OVERLAP",\n'
|
||||||
|
' "severity": "warning",\n'
|
||||||
|
' "message": "Описание проблемы на русском",\n'
|
||||||
|
' "bbox": [[x1,y1],[x2,y2],[x3,y3],[x4,y4]]\n'
|
||||||
|
' }\n'
|
||||||
|
']\n\n'
|
||||||
|
"Если проблем нет — верни пустой массив [].\n"
|
||||||
|
"severity: 'error' (критично), 'warning' (стоит исправить), 'info' (замечание).\n"
|
||||||
|
"bbox — координаты проблемной зоны в пикселях (если можешь определить),"
|
||||||
|
" иначе верни null."
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
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 parse_vlm_response(text: str) -> list:
|
||||||
|
"""Парсит 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:
|
||||||
|
data = json.loads(text)
|
||||||
|
if isinstance(data, list):
|
||||||
|
return data
|
||||||
|
elif isinstance(data, dict) and "issues" in data:
|
||||||
|
return data["issues"]
|
||||||
|
else:
|
||||||
|
return []
|
||||||
|
except json.JSONDecodeError as e:
|
||||||
|
print(f"[WARN] Не удалось распарсить JSON: {e}")
|
||||||
|
print(f"[WARN] Raw text: {text[:500]}")
|
||||||
|
return []
|
||||||
|
|
||||||
|
|
||||||
|
def analyze_page(image_path: Path, model: str) -> list:
|
||||||
|
"""Отправляет PNG в qwen-vl API, получает список issues."""
|
||||||
|
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": QC_PROMPT},
|
||||||
|
{"type": "image_url", "image_url": {"url": data_url}},
|
||||||
|
],
|
||||||
|
}
|
||||||
|
],
|
||||||
|
temperature=0.2,
|
||||||
|
max_tokens=4096,
|
||||||
|
)
|
||||||
|
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")
|
||||||
|
|
||||||
|
issues = parse_vlm_response(raw)
|
||||||
|
|
||||||
|
# Масштабировать bbox обратно к оригиналу
|
||||||
|
if scale != 1.0:
|
||||||
|
for issue in issues:
|
||||||
|
bbox = issue.get("bbox")
|
||||||
|
if bbox and isinstance(bbox, list):
|
||||||
|
for point in bbox:
|
||||||
|
if isinstance(point, list) and len(point) == 2:
|
||||||
|
point[0] = round(point[0] / scale)
|
||||||
|
point[1] = round(point[1] / scale)
|
||||||
|
|
||||||
|
return issues
|
||||||
|
|
||||||
|
|
||||||
|
def run_vlm_qc(folder: Path, model: str = DEFAULT_MODEL):
|
||||||
|
"""Запускает VLM-QC для всех 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_qc_report.json"
|
||||||
|
report = {"errors": [], "warnings": [], "infos": [], "source": "vlm"}
|
||||||
|
|
||||||
|
print(f"[INFO] VLM QC: {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:
|
||||||
|
issues = analyze_page(png, model)
|
||||||
|
|
||||||
|
page_num = int(png.stem.split("_")[1])
|
||||||
|
for issue in issues:
|
||||||
|
issue["page"] = page_num
|
||||||
|
issue["source"] = "vlm"
|
||||||
|
sev = issue.get("severity", "warning")
|
||||||
|
if sev not in ("error", "warning", "info"):
|
||||||
|
sev = "warning"
|
||||||
|
report[f"{sev}s"].append(issue)
|
||||||
|
|
||||||
|
print(f"OK ({len(issues)} issues)")
|
||||||
|
except Exception as e:
|
||||||
|
print(f"ERR: {e}")
|
||||||
|
|
||||||
|
with open(out_path, "w", encoding="utf-8") as f:
|
||||||
|
json.dump(report, f, ensure_ascii=False, indent=2)
|
||||||
|
|
||||||
|
total = sum(len(report[k]) for k in ["errors", "warnings", "infos"])
|
||||||
|
print(f"\n[OK] VLM QC сохранён: {out_path}")
|
||||||
|
print(f" Всего замечаний: {total}")
|
||||||
|
print(f" Ошибки: {len(report['errors'])}, Предупреждения: {len(report['warnings'])}, Инфо: {len(report['infos'])})")
|
||||||
|
|
||||||
|
|
||||||
|
def main():
|
||||||
|
import argparse
|
||||||
|
parser = argparse.ArgumentParser(description="VLM QC для чертежей через qwen-vl API")
|
||||||
|
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_qc(folder, args.model)
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
main()
|
||||||
Loading…
Reference in New Issue
Block a user