HRTx / inference.py
d0r1h
add app files
83c00db
raw
history blame contribute delete
480 Bytes
from models.trocr import TrOCROCR
# Lazy-loaded registry (important for big models)
_MODEL_CACHE = {}
def get_model(model_name: str):
if model_name not in _MODEL_CACHE:
if model_name == "trocr":
_MODEL_CACHE[model_name] = TrOCROCR()
else:
raise ValueError(f"Unknown model: {model_name}")
return _MODEL_CACHE[model_name]
def run_ocr(image, model_name="trocr"):
model = get_model(model_name)
return model.predict(image)