File size: 1,029 Bytes
479f206 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 | import os
import requests
HF_MODEL = os.getenv("HF_OCR_MODEL", "microsoft/trocr-base-printed")
def run_hf_ocr(image_bytes: bytes) -> str:
token = os.getenv("HUGGINGFACE_API_TOKEN")
if not token:
raise RuntimeError("Missing HUGGINGFACE_API_TOKEN")
# Use Inference API for text recognition
url = f"https://api-inference.huggingface.co/models/{HF_MODEL}"
headers = {"Authorization": f"Bearer {token}"}
response = requests.post(url, headers=headers, data=image_bytes, timeout=60)
if response.status_code >= 400:
raise RuntimeError(f"HF API error: {response.status_code} {response.text}")
try:
data = response.json()
except Exception as exc:
raise RuntimeError(f"Invalid HF response: {exc}")
# trocr returns dicts with 'generated_text'
if isinstance(data, list) and data and isinstance(data[0], dict):
text = data[0].get("generated_text")
if isinstance(text, str):
return text
# fallback to raw
return str(data)
|