54 lines
1.3 KiB
Python
54 lines
1.3 KiB
Python
import os
|
|
import time
|
|
import requests
|
|
from dotenv import load_dotenv
|
|
|
|
load_dotenv()
|
|
|
|
API_URL = os.getenv('OCR_API_URL')
|
|
HEADERS = {
|
|
"Authorization": f"bearer {os.getenv('OCR_API_TOKEN')}",
|
|
"Content-Type": "application/json",
|
|
}
|
|
MODEL = os.getenv("OCR_MODEL")
|
|
|
|
|
|
def ocr_from_url(image_url: str) -> str:
|
|
r = requests.post(
|
|
API_URL,
|
|
json={
|
|
"fileUrl": image_url,
|
|
"model": MODEL,
|
|
"optionalPayload": {
|
|
"useDocOrientationClassify": False,
|
|
"useDocUnwarping": False,
|
|
"useChartRecognition": False,
|
|
},
|
|
},
|
|
headers=HEADERS,
|
|
)
|
|
r.raise_for_status()
|
|
job_id = r.json()["data"]["jobId"]
|
|
|
|
while True:
|
|
r = requests.get(f"{API_URL}/{job_id}", headers=HEADERS)
|
|
r.raise_for_status()
|
|
data = r.json()["data"]
|
|
|
|
if data["state"] == "done":
|
|
jsonl_url = data["resultUrl"]["jsonUrl"]
|
|
break
|
|
if data["state"] == "failed":
|
|
raise RuntimeError("OCR failed")
|
|
|
|
time.sleep(3)
|
|
|
|
lines = requests.get(jsonl_url).text.strip().splitlines()
|
|
text = ""
|
|
for line in lines:
|
|
item = __import__("json").loads(line)
|
|
for res in item["result"]["layoutParsingResults"]:
|
|
text += res["markdown"]["text"]
|
|
|
|
return text
|