235 lines
8.9 KiB
Python
235 lines
8.9 KiB
Python
|
#!/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()
|