tekBlueprint

#!/usr/bin/env bash

LLAMA_SERVER=path/to/llama-server
MODEL=path/to/model/gemma-4-12b-it-UD-Q8_K_XL.gguf
MODEL_MTP=path/to/model/mtp-gemma-4-12b-it.gguf
MODEL_MM=path/to/model/mmproj-F32.gguf

$LLAMA_SERVER \
    --host 0.0.0.0 \
    --port 1234 \
    -c 0 \
    --batch-size 2048 \
    --ubatch-size 2048 \
    --temp 0.1 \
    --top-p 0.9 \
    --top-k 40 \
    --min-p 0.01 \
    --repeat-penalty 1.0 \
    --presence-penalty 0.0 \
    -fa 1 \
    -ngl -1 \
    --image-min-tokens 1120 \
    --image-max-tokens 1120 \
    -t 16 \
    -tb 32 \
    --prio 2 \
    --cpu-strict 1 \
    --no-mmap \
    --no-warmup \
    --jinja \
    --alias "gemma-4-12b-it-ocr-mtp-q8_k_xl" \
    --spec-type draft-mtp \
    --spec-draft-model ${MODEL_MTP} \
    --spec-draft-n-max 3 \
    --mmproj ${MODEL_MM} \
    -m ${MODEL}

import os
import time
import base64
import argparse
import subprocess
import requests
from concurrent.futures import ThreadPoolExecutor, as_completed

API_URL_DEFAULT = "http://localhost:1234/v1/chat/completions"
MODEL_NAME_DEFAULT = "gemma-4-12b-it-ocr-mtp-q8_k_xl"
SYSTEM_PROMPT = """You are a precise OCR assistant. Transcribe all visible text from the image into clean Markdown.

Rules:
- Use `##` headings for document sections and subsections.
- Bold field labels with `**Label**: value` syntax where applicable.
- Separate logical sections with blank lines.
- Preserve tables, lists, and formatting as close to the original as possible.
- Keep all numbers, units, dates, and special characters exactly as written."""

SKIP_DIRS = {'nemotron_env', 'nemotron_src', '__pycache__', '.git', 'old'}

SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
IMAGES_DIR = os.path.join(SCRIPT_DIR, 'images')
DOCS_DIR = os.path.join(SCRIPT_DIR, 'docs')


def prepare_image(image_path, max_dimension=None, sharpen=False):
    ext = os.path.splitext(image_path)[1].lower()
    cleanup_files = []

    if ext == '.webp':
        png_path = image_path.replace('.webp', '.png')
        try:
            subprocess.run(['ffmpeg', '-y', '-i', image_path, png_path, '-loglevel', 'error'], check=True)
        except Exception as e:
            return None, f"Error converting {image_path}: {e}", []
        work_path = png_path
        cleanup_files.append(work_path)
    elif ext in ['.png', '.jpg', '.jpeg']:
        work_path = image_path
    else:
        return None, f"Error: Unsupported format {ext}", []

    if sharpen or max_dimension:
        preprocessed = work_path + '.temp_preprocessed.png'
        try:
            cmd = ['magick', work_path]
            if sharpen:
                cmd.extend(['-unsharp', '0x2+1+0'])
            if max_dimension:
                cmd.extend([
                    '-resize',
                    f'{max_dimension}x{max_dimension}>'
                ])
            cmd.append(preprocessed)
            subprocess.run(cmd, check=True)
        except Exception as e:
            for f in cleanup_files:
                os.unlink(f)
            return None, f"Error preprocessing image: {e}", []
        if work_path == image_path:
            cleanup_files.append(preprocessed)
        else:
            work_path = preprocessed
            cleanup_files[-1] = preprocessed

    return work_path, None, cleanup_files


def encode_image(image_path):
    with open(image_path, "rb") as f:
        return base64.b64encode(f.read()).decode('utf-8')


def recognize_image(image_path, api_url, model_name, timeout, max_dimension, sharpen, max_retries=3):
    work_path, error, cleanup_files = prepare_image(image_path, max_dimension, sharpen)
    if error:
        return error

    ext = os.path.splitext(work_path)[1].lower()
    mime_type = {".png": "image/png", ".jpg": "image/jpeg", ".jpeg": "image/jpeg"}.get(ext, "image/png")
    base64_image = encode_image(work_path)

    payload = {
        "model": model_name,
        "messages": [
            {"role": "system", "content": SYSTEM_PROMPT},
            {
                "role": "user",
                "content": [
                    {"type": "text", "text": "Transcribe all visible text from this document into structured Markdown."},
                    {"type": "image_url", "image_url": {"url": f"data:{mime_type};base64,{base64_image}"}}
                ]
            }
        ],
        "stream": False,
        "temperature": 0.1,
        "max_tokens": 16384,
        "chat_template_kwargs": {"enable_thinking": False}
    }

    last_error = None
    for attempt in range(1, max_retries + 1):
        try:
            response = requests.post(api_url, json=payload, timeout=timeout)
            response.raise_for_status()
            result = response.json()
            content = result['choices'][0]['message']['content'].strip()

            for f in cleanup_files:
                if os.path.exists(f):
                    os.remove(f)

            if not content:
                return "Warning: Model returned an empty transcription."
            return content

        except requests.exceptions.RequestException as e:
            last_error = e
            if attempt < max_retries:
                backoff = 2 ** (attempt - 1)
                print(f"    Retry {attempt}/{max_retries} after error ({type(e).__name__}): sleeping {backoff}s...")
                time.sleep(backoff)

    for f in cleanup_files:
        if os.path.exists(f):
            os.remove(f)

    return f"Error calling API after {max_retries} attempts: {last_error}"


def process_file(file_path, api_url, model_name, timeout, max_dimension, sharpen, force):
    base_name = os.path.splitext(os.path.basename(file_path))[0]
    md_file = os.path.join(DOCS_DIR, f"{base_name}.md")

    if not force and os.path.exists(md_file) and os.path.getsize(md_file) > 10:
        print(f"[skip] {file_path} (MD already exists)")
        return

    print(f"[ocr]  {file_path}")
    text_content = recognize_image(file_path, api_url, model_name, timeout, max_dimension, sharpen)

    if text_content.startswith("Error") or text_content.startswith("Warning"):
        print(f"       Failed: {text_content}")
    else:
        os.makedirs(DOCS_DIR, exist_ok=True)
        with open(md_file, 'w', encoding='utf-8') as f:
            f.write(f"# {base_name}\n\n")
            f.write(text_content)
        print(f"       Saved to {md_file}")


def worker(file_path, api_url, model_name, timeout, max_dimension, sharpen, force):
    process_file(file_path, api_url, model_name, timeout, max_dimension, sharpen, force)


def main():
    parser = argparse.ArgumentParser(
        description="OCR images to Markdown using Vision LLM API (OpenAI-compatible)"
    )
    parser.add_argument('--api-url', type=str, default=API_URL_DEFAULT,
                        help=f'API endpoint URL (default: {API_URL_DEFAULT})')
    parser.add_argument('--model', type=str, default=MODEL_NAME_DEFAULT,
                        help=f'Model name (default: {MODEL_NAME_DEFAULT})')
    parser.add_argument('--timeout', type=int, default=1200,
                        help='Request timeout in seconds (default: 1200)')
    parser.add_argument('--max-dimension', type=int, default=None,
                        help='Resize images to max dimension on long side (default: no resize)')
    parser.add_argument('--sharpen', action='store_true',
                        help='Apply unsharp mask before sending to API')
    parser.add_argument('--force', action='store_true',
                        help='Reprocess even if .md already exists')
    parser.add_argument('--workers', type=int, default=1,
                        help='Number of parallel workers (default: 1)')
    args, remaining = parser.parse_known_args()

    target_file = remaining[0] if remaining else None

    if target_file:
        if os.path.exists(target_file):
            process_file(
                target_file, args.api_url, args.model,
                args.timeout, args.max_dimension, args.sharpen, args.force
            )
        else:
            print(f"Error: File '{target_file}' not found.")
        return

    if not os.path.isdir(IMAGES_DIR):
        print(f"Error: Images directory '{IMAGES_DIR}' not found.")
        return

    valid_extensions = ('.webp', '.png', '.jpg', '.jpeg')
    files_to_process = []
    for root, dirs, files in os.walk(IMAGES_DIR):
        dirs[:] = [d for d in dirs if d not in SKIP_DIRS]
        for f in files:
            if f.lower().endswith(valid_extensions):
                files_to_process.append(os.path.join(root, f))

    if not files_to_process:
        print("No valid image files found.")
        return

    preprocess_info = []
    if args.max_dimension:
        preprocess_info.append(f"resize={args.max_dimension}")
    if args.sharpen:
        preprocess_info.append("sharpen")
    info_str = f", {'+'.join(preprocess_info)}" if preprocess_info else ""
    print(f"Model: {args.model}, workers: {args.workers}{info_str}\n")

    if args.workers == 1:
        for file_path in files_to_process:
            process_file(
                file_path, args.api_url, args.model,
                args.timeout, args.max_dimension, args.sharpen, args.force
            )
    else:
        with ThreadPoolExecutor(max_workers=args.workers) as executor:
            futures = {
                executor.submit(
                    worker, fp, args.api_url, args.model,
                    args.timeout, args.max_dimension, args.sharpen, args.force
                ): fp for fp in files_to_process
            }
            for future in as_completed(futures):
                try:
                    future.result()
                except Exception as e:
                    print(f"[error] {futures[future]}: {e}")


if __name__ == "__main__":
    main()

python doc_to_md_vision.py /path/to/document.png

# Последовательно (по умолчанию)
python doc_to_md_vision.py

# Параллельно, 2 воркера
python doc_to_md_vision.py --workers 2

# Улучшить резкость + ограничить размер
python doc_to_md_vision.py --sharpen --max-dimension 2048

# Пересоздать существующие .md
python doc_to_md_vision.py --force