convert.py

"""End-to-end: download eDOCr2 weights and convert all three stages to CoreML.

Usage:
    pip install -r requirements.txt
    python convert.py

Outputs:
    artefacts/edocr2_detector.mlpackage          (CRAFT text detector)
    artefacts/edocr2_recogniser.mlpackage        (CRNN dimension recogniser)
    artefacts/edocr2_gdt_classifier.mlpackage    (CRNN GD&T recogniser)

The upstream eDOCr2 source is cloned into ./edocr2-upstream/ so its
Keras module definitions can be imported during conversion.
"""
import hashlib
import os
import shutil
import subprocess
import sys
import time
import urllib.request

os.environ.setdefault("TF_USE_LEGACY_KERAS", "1")
os.environ.setdefault("TF_CPP_MIN_LOG_LEVEL", "3")

ROOT = os.path.dirname(os.path.abspath(__file__))
UPSTREAM_DIR = os.path.join(ROOT, "edocr2-upstream")
UPSTREAM_URL = "https://github.com/javvi51/edocr2.git"
UPSTREAM_COMMIT = "f6f96517a531021ac946f6fc45063bdb77440085"  # main @ 2025-01-22

WEIGHTS_DIR = os.path.join(ROOT, "weights")
ARTEFACT_DIR = os.path.join(ROOT, "artefacts")

WEIGHTS = {
    # CRAFT detector, via the keras-ocr v0.8.4 release (upstream-compatible).
    "craft_mlt_25k.h5": {
        "url": "https://github.com/faustomorales/keras-ocr/releases/download/v0.8.4/craft_mlt_25k.h5",
        "sha256": "7283ce2ff05a0617e9740c316175ff3bacdd7215dbdf1a726890d5099431f899",
    },
    # Trained recognisers published with eDOCr2 v1.0.0.
    "recognizer_dimensions_2.keras": {
        "url": "https://github.com/javvi51/edocr2/releases/download/v1.0.0/recognizer_dimensions_2.keras",
    },
    "recognizer_dimensions_2.txt": {
        "url": "https://github.com/javvi51/edocr2/releases/download/v1.0.0/recognizer_dimensions_2.txt",
    },
    "recognizer_gdts.keras": {
        "url": "https://github.com/javvi51/edocr2/releases/download/v1.0.0/recognizer_gdts.keras",
    },
    "recognizer_gdts.txt": {
        "url": "https://github.com/javvi51/edocr2/releases/download/v1.0.0/recognizer_gdts.txt",
    },
}

DETECTOR_H = 1280
DETECTOR_W = 1280


def download(name: str, info: dict) -> str:
    path = os.path.join(WEIGHTS_DIR, name)
    if os.path.exists(path):
        return path
    os.makedirs(WEIGHTS_DIR, exist_ok=True)
    print(f"  downloading {name}...")
    urllib.request.urlretrieve(info["url"], path)
    if "sha256" in info:
        h = hashlib.sha256(open(path, "rb").read()).hexdigest()
        assert h == info["sha256"], f"sha256 mismatch for {name}: got {h}"
    return path


def clone_upstream():
    if os.path.isdir(UPSTREAM_DIR) and os.path.isdir(os.path.join(UPSTREAM_DIR, ".git")):
        return
    print(f"Cloning upstream eDOCr2 ({UPSTREAM_COMMIT[:8]})...")
    subprocess.check_call(["git", "clone", UPSTREAM_URL, UPSTREAM_DIR])
    subprocess.check_call(["git", "-C", UPSTREAM_DIR, "checkout", UPSTREAM_COMMIT])


def _save_mlpackage(mlmodel, dest: str):
    if os.path.isdir(dest):
        shutil.rmtree(dest)
    os.makedirs(os.path.dirname(dest), exist_ok=True)
    mlmodel.save(dest)


def convert_detector():
    import numpy as np
    import tensorflow as tf
    import coremltools as ct
    import tf2onnx

    sys.path.insert(0, UPSTREAM_DIR)
    from edocr2.keras_ocr import detection

    from tensorflow import keras
    original_input = keras.layers.Input
    def fixed_input(shape=(None, None, 3), **kw):
        return original_input(shape=(DETECTOR_H, DETECTOR_W, 3), batch_size=1,
                              **{k: v for k, v in kw.items() if k != "batch_size"})
    keras.layers.Input = fixed_input
    try:
        model = detection.build_keras_model(
            weights_path=os.path.join(WEIGHTS_DIR, "craft_mlt_25k.h5"),
            backbone_name="vgg",
        )
    finally:
        keras.layers.Input = original_input

    print(f"  detector params={model.count_params():,} input={model.input_shape} output={model.output_shape}")

    spec = (tf.TensorSpec((1, DETECTOR_H, DETECTOR_W, 3), tf.float32, name="input"),)
    onnx_path = os.path.join(WEIGHTS_DIR, "edocr2_detector.onnx")
    tf2onnx.convert.from_keras(model, input_signature=spec, opset=17, output_path=onnx_path)

    tf_input_name = model.inputs[0].name.split(":")[0]
    mlmodel = ct.convert(
        model,
        source="tensorflow",
        inputs=[ct.TensorType(name=tf_input_name, shape=(1, DETECTOR_H, DETECTOR_W, 3))],
        convert_to="mlprogram",
        compute_precision=ct.precision.FLOAT16,
        compute_units=ct.ComputeUnit.ALL,
        minimum_deployment_target=ct.target.macOS15,
    )
    mlmodel.short_description = (
        f"eDOCr2 CRAFT text detector (VGG backbone). Input: NHWC RGB {DETECTOR_H}x{DETECTOR_W} "
        "normalised with ImageNet mean/std. Output: two-channel heatmap (region, affinity)."
    )
    mlmodel.author = "Converted from javvi51/edocr2 (MIT) via tf2onnx + coremltools"
    mlmodel.user_defined_metadata["input_height"] = str(DETECTOR_H)
    mlmodel.user_defined_metadata["input_width"] = str(DETECTOR_W)
    mlmodel.user_defined_metadata["normalisation"] = "imagenet (mean=[0.485,0.456,0.406], std=[0.229,0.224,0.225])"
    _save_mlpackage(mlmodel, os.path.join(ARTEFACT_DIR, "edocr2_detector.mlpackage"))


def convert_recogniser(variant: str):
    import numpy as np
    import tensorflow as tf
    import coremltools as ct
    import tf2onnx

    sys.path.insert(0, UPSTREAM_DIR)
    from edocr2.keras_ocr.recognition import Recognizer

    mapping = {
        "dimensions": ("recognizer_dimensions_2.keras", "recognizer_dimensions_2.txt", "edocr2_recogniser.mlpackage",
                       "eDOCr2 CRNN recogniser for dimensions (digits + A-Za-z + dim symbols)."),
        "gdts": ("recognizer_gdts.keras", "recognizer_gdts.txt", "edocr2_gdt_classifier.mlpackage",
                 "eDOCr2 CRNN recogniser for GD&T symbols (∅⌖⌒⌓⏤⏥⏊⌭⫽◎↗⌰⌯ + datum letters)."),
    }
    weights_file, alphabet_file, mlpkg_name, description = mapping[variant]

    weights_path = os.path.join(WEIGHTS_DIR, weights_file)
    alphabet = open(os.path.join(WEIGHTS_DIR, alphabet_file)).read()

    recognizer = Recognizer(alphabet=alphabet)
    recognizer.model.load_weights(weights_path)
    model = recognizer.model
    _, h, w, c = model.input_shape
    print(f"  {variant} alphabet={len(alphabet)} input={model.input_shape} output={model.output_shape}")

    spec = (tf.TensorSpec((1, h, w, c), tf.float32, name="image"),)
    onnx_path = os.path.join(WEIGHTS_DIR, mlpkg_name.replace(".mlpackage", ".onnx"))
    tf2onnx.convert.from_keras(model, input_signature=spec, opset=17, output_path=onnx_path)

    tf_input_name = model.inputs[0].name.split(":")[0]
    mlmodel = ct.convert(
        model,
        source="tensorflow",
        inputs=[ct.TensorType(name=tf_input_name, shape=(1, h, w, c))],
        convert_to="mlprogram",
        compute_precision=ct.precision.FLOAT16,
        compute_units=ct.ComputeUnit.ALL,
        minimum_deployment_target=ct.target.macOS15,
    )
    mlmodel.short_description = description
    mlmodel.author = "Converted from javvi51/edocr2 (MIT) via coremltools"
    mlmodel.user_defined_metadata["alphabet"] = alphabet
    mlmodel.user_defined_metadata["blank_index"] = str(len(alphabet))
    mlmodel.user_defined_metadata["input_height"] = str(h)
    mlmodel.user_defined_metadata["input_width"] = str(w)
    mlmodel.user_defined_metadata["channels"] = str(c)
    _save_mlpackage(mlmodel, os.path.join(ARTEFACT_DIR, mlpkg_name))


def main():
    os.makedirs(WEIGHTS_DIR, exist_ok=True)
    for name, info in WEIGHTS.items():
        download(name, info)

    clone_upstream()

    print("Converting CRAFT detector...")
    t0 = time.time()
    convert_detector()
    print(f"  done in {time.time()-t0:.1f}s")

    print("Converting dimension recogniser...")
    t0 = time.time()
    convert_recogniser("dimensions")
    print(f"  done in {time.time()-t0:.1f}s")

    print("Converting GD&T recogniser...")
    t0 = time.time()
    convert_recogniser("gdts")
    print(f"  done in {time.time()-t0:.1f}s")

    print("\nAll three CoreML packages written to artefacts/.")


if __name__ == "__main__":
    main()