Blame for install_ocr_models.py

"""Install pinned, checksum-verified Chinese OCR models while building the image."""
import hashlib
from pathlib import Path
import re
import subprocess
import urllib.request

REVISION = 'e12c65a915945e4c28e237a9b52bc4a8f39a0cec'
MODELS = {
    'chi_sim': '4fef2d1306c8e87616d4d3e4c6c67faf5d44be3342290cf8f2f0f6e3aa7e735b',
    'chi_tra': '1aa60488574cafa69486d919284f079ca9b68fcc7f6ad8dc1ff1b318dfd97028',
}


def install():
    result = subprocess.run(['tesseract', '--list-langs'], capture_output=True, text=True, check=True)
    match = re.search(r'in "([^"]+)"', result.stdout + result.stderr)
    if not match or not Path(match[1]).is_dir():
        raise RuntimeError('Cannot locate the installed Tesseract data directory.')
    directory = Path(match[1])
    for language, expected in MODELS.items():
        url = f'https://raw.githubusercontent.com/tesseract-ocr/tessdata_best/{REVISION}/{language}.traineddata'
        with urllib.request.urlopen(url, timeout=60) as stream:
            data = stream.read()
        if hashlib.sha256(data).hexdigest() != expected:
            raise RuntimeError(f'{language}: model checksum mismatch')
        (directory / (language + '.traineddata')).write_bytes(data)
        print(f'Installed verified {language} model at {REVISION}')


if __name__ == '__main__':
    install()