You cannot select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
81 lines
2.2 KiB
Python
81 lines
2.2 KiB
Python
from fastapi import FastAPI
|
|
from pydantic import BaseModel
|
|
import base64
|
|
import subprocess
|
|
import re
|
|
import uuid
|
|
import os
|
|
import logging
|
|
from typing import List
|
|
|
|
app = FastAPI()
|
|
|
|
logging.basicConfig(level=logging.INFO)
|
|
logger = logging.getLogger(__name__)
|
|
|
|
ansi_escape = re.compile(r'\x1b[^m]*m')
|
|
|
|
|
|
class PictureRequest(BaseModel):
|
|
pic_id: int
|
|
pic: str
|
|
|
|
|
|
class PictureResponse(BaseModel):
|
|
pic_id: int
|
|
pic_txt: str
|
|
code: int
|
|
|
|
|
|
@app.post("/ocr", response_model=List[PictureResponse])
|
|
async def ocr_endpoint(pictures: List[PictureRequest]):
|
|
results = []
|
|
|
|
for picture in pictures:
|
|
image_data = base64.b64decode(picture.pic)
|
|
temp_image_file = f"{uuid.uuid4()}.jpg" # 使用uuid生成唯一的临时文件名
|
|
|
|
try:
|
|
with open(temp_image_file, "wb") as f:
|
|
f.write(image_data)
|
|
|
|
command = [
|
|
'python', 'tools/infer/predict_system_1.py',
|
|
'--use_gpu=False',
|
|
'--cls_model_dir=./models/cls',
|
|
'--rec_model_dir=./models/rec',
|
|
'--det_model_dir=./models/det',
|
|
f'--image_dir={temp_image_file}'
|
|
]
|
|
|
|
logger.info(f"Processing image with pic_id: {picture.pic_id}")
|
|
result = subprocess.run(command, capture_output=True, text=True, encoding='utf-8', errors='ignore')
|
|
|
|
if result.returncode == 0:
|
|
ocr_str = ansi_escape.sub('', result.stdout).strip()
|
|
code = 200
|
|
logger.info(f"OCR successful for pic_id: {picture.pic_id}")
|
|
else:
|
|
ocr_str = ""
|
|
code = 500
|
|
logger.error(f"OCR failed for pic_id: {picture.pic_id}, return code: {result.returncode}")
|
|
|
|
except Exception as e:
|
|
ocr_str = str(e)
|
|
code = 500
|
|
logger.exception(f"Exception occurred while processing pic_id: {picture.pic_id}")
|
|
|
|
finally:
|
|
if os.path.exists(temp_image_file):
|
|
os.remove(temp_image_file) # 确保临时文件在处理完成后被删除
|
|
|
|
results.append({"pic_id": picture.pic_id, "pic_txt": ocr_str, "code": code})
|
|
|
|
return results
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import uvicorn
|
|
|
|
uvicorn.run(app, host="0.0.0.0", port=8000)
|