File size: 2,324 Bytes
c50f0de | 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 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 | """
HuggingFace Inference API client for Donut OCR.
Use this instead of local model loading for lightweight cloud deployment.
"""
import base64
import io
import json
import requests
from PIL import Image
HF_API_URL = "https://api-inference.huggingface.co/models"
def run_donut_hf_api(image, model_id="naver-clova-ix/donut-base-finetuned-cord-v2", api_token=None):
"""
Run Donut OCR via HuggingFace Inference API.
Args:
image: PIL Image or path to image
model_id: HuggingFace model ID (e.g., "your-username/your-model")
api_token: HuggingFace API token (get from https://huggingface.co/settings/tokens)
Returns:
dict with "ocr_engine", "raw_sequence", "parsed"
"""
if isinstance(image, str):
image = Image.open(image).convert("RGB")
elif not isinstance(image, Image.Image):
raise ValueError("image must be PIL Image or file path")
# Convert to bytes
buf = io.BytesIO()
image.save(buf, format="PNG")
img_bytes = buf.getvalue()
# Call HF Inference API
headers = {}
if api_token:
headers["Authorization"] = f"Bearer {api_token}"
api_endpoint = f"{HF_API_URL}/{model_id}"
try:
response = requests.post(
api_endpoint,
headers=headers,
data=img_bytes,
timeout=60
)
response.raise_for_status()
result = response.json()
# HF API returns a list, get first element
if isinstance(result, list):
result = result[0]
# Extract generated text
generated_text = result.get("generated_text", "")
return {
"ocr_engine": f"HF API ({model_id})",
"raw_sequence": generated_text,
"parsed": None # You'd need to parse the sequence same as local Donut
}
except requests.exceptions.HTTPError as e:
if response.status_code == 503:
return {"error": "Model is loading on HF. Wait a minute and retry."}
elif response.status_code == 401:
return {"error": "Invalid HF API token. Get one at huggingface.co/settings/tokens"}
else:
return {"error": f"HF API error {response.status_code}: {response.text}"}
except Exception as e:
return {"error": f"Request failed: {str(e)}"}
|