c756a5766b
- Add rag_indexer.py: build LightRAG index from OCR with OpenCode API - Add rag_query.py: query the knowledge graph - Add vlm_describer.py: generate VLM descriptions via LM Studio - Add test_model.py: quick check for LightRAG-compatible models - Add run_pipeline.sh and run_pipeline.bat: full OCR → VLM → RAG pipeline - Fix rapidocr import (rapidocr_onnxruntime) - Fix process_any_pdf.py paths for cross-platform use - Add .env.example, README_RAG.md, AGENTS.md - Update .gitignore for outputs and secrets
115 lines
4.0 KiB
Python
115 lines
4.0 KiB
Python
#!/usr/bin/env python3
|
||
# -*- coding: utf-8 -*-
|
||
"""
|
||
Генерация текстовых описаний PNG-страниц через VLM в LM Studio.
|
||
|
||
Требования:
|
||
- Запущен LM Studio с загруженной моделью (например, qwen3-vl-4b)
|
||
- Сервер: http://127.0.0.1:1234/v1
|
||
|
||
Использование:
|
||
python vlm_describer.py <output_folder> [--prompt "..."] [--model MODEL]
|
||
|
||
Результат: <output_folder>/vlm_descriptions.json
|
||
"""
|
||
|
||
import os
|
||
import sys
|
||
import json
|
||
import base64
|
||
import argparse
|
||
from pathlib import Path
|
||
from openai import OpenAI
|
||
|
||
# ------------------------------------------------------------------
|
||
# Конфигурация LM Studio
|
||
# ------------------------------------------------------------------
|
||
LMSTUDIO_URL = os.environ.get("LMSTUDIO_URL", "http://127.0.0.1:1234/v1")
|
||
LMSTUDIO_KEY = os.environ.get("LMSTUDIO_API_KEY", "lm-studio")
|
||
|
||
DEFAULT_PROMPT = (
|
||
"Опиши этот чертеж подробно. Укажи:\n"
|
||
"- Какой это этаж (если видно)\n"
|
||
"- Какие оси обозначены\n"
|
||
"- Какие размеры указаны\n"
|
||
"- Какие помещения/квартиры видны\n"
|
||
"- Общую компоновку и заметные детали.\n"
|
||
"Отвечай на русском языке."
|
||
)
|
||
|
||
client = OpenAI(base_url=LMSTUDIO_URL, api_key=LMSTUDIO_KEY)
|
||
|
||
|
||
def encode_image(image_path: Path) -> str:
|
||
with open(image_path, "rb") as f:
|
||
return base64.b64encode(f.read()).decode("utf-8")
|
||
|
||
|
||
def describe_image(image_path: Path, model: str, prompt: str) -> str:
|
||
"""Отправляет PNG в VLM и получает текстовое описание."""
|
||
b64 = encode_image(image_path)
|
||
data_url = f"data:image/png;base64,{b64}"
|
||
|
||
response = client.chat.completions.create(
|
||
model=model,
|
||
messages=[
|
||
{
|
||
"role": "user",
|
||
"content": [
|
||
{"type": "text", "text": prompt},
|
||
{"type": "image_url", "image_url": {"url": data_url}},
|
||
],
|
||
}
|
||
],
|
||
temperature=0.3,
|
||
max_tokens=512, # 4B модель быстро устаёт, не гоним длину
|
||
)
|
||
return response.choices[0].message.content.strip()
|
||
|
||
|
||
def process_folder(folder: Path, model: str, prompt: str):
|
||
"""Обрабатывает все 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_descriptions.json"
|
||
descriptions = {}
|
||
|
||
print(f"[INFO] Найдено {len(png_files)} изображений")
|
||
print(f"[INFO] LM Studio: {LMSTUDIO_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:
|
||
desc = describe_image(png, model, prompt)
|
||
descriptions[png.name] = desc
|
||
print(f"OK ({len(desc)} chars)")
|
||
except Exception as e:
|
||
print(f"ERR: {e}")
|
||
descriptions[png.name] = f"[ERROR] {e}"
|
||
|
||
with open(out_path, "w", encoding="utf-8") as f:
|
||
json.dump(descriptions, f, ensure_ascii=False, indent=2)
|
||
|
||
print(f"\n[OK] Сохранено: {out_path}")
|
||
|
||
|
||
def main():
|
||
parser = argparse.ArgumentParser(description="VLM-описания PNG через LM Studio")
|
||
parser.add_argument("folder", help="Папка с page_*.png")
|
||
parser.add_argument("--model", default="qwen/qwen3-vl-4b",
|
||
help="Имя модели в LM Studio (default: qwen/qwen3-vl-4b)")
|
||
parser.add_argument("--prompt", default=DEFAULT_PROMPT,
|
||
help="Промпт для VLM")
|
||
args = parser.parse_args()
|
||
|
||
folder = Path(args.folder)
|
||
process_folder(folder, args.model, args.prompt)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|