File size: 480 Bytes
83c00db | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 | 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)
|