wehub-resource-sync
59 lines
1.4 KiB
Python
59 lines
1.4 KiB
Python
import base64
|
|
import gc
|
|
import io
|
|
|
|
import numpy as np
|
|
from fastapi import FastAPI
|
|
from paddleocr import PaddleOCR
|
|
from PIL import Image
|
|
from pydantic import BaseModel
|
|
|
|
app = FastAPI()
|
|
ocr_module = PaddleOCR(use_angle_cls=True, lang="en")
|
|
|
|
|
|
class ImageData(BaseModel):
|
|
img_bytes: bytes
|
|
|
|
|
|
def text_cvt_orc_format_paddle(paddle_result):
|
|
texts = []
|
|
print("paddle_result: ", paddle_result)
|
|
for i, line in enumerate(paddle_result[0]):
|
|
points = np.array(line[0])
|
|
print("points: ", points)
|
|
location = {
|
|
"left": int(min(points[:, 0])),
|
|
"top": int(min(points[:, 1])),
|
|
"right": int(max(points[:, 0])),
|
|
"bottom": int(max(points[:, 1])),
|
|
}
|
|
print("location: ", location)
|
|
content = line[1][0]
|
|
texts.append((i, content, location))
|
|
return texts
|
|
|
|
|
|
def ocr_results(screenshot):
|
|
screenshot_img = Image.open(io.BytesIO(screenshot))
|
|
result = ocr_module.ocr(np.array(screenshot_img), cls=True)
|
|
return text_cvt_orc_format_paddle(result)
|
|
|
|
|
|
@app.post("/ocr/")
|
|
async def read_image(image_data: ImageData):
|
|
image_bytes = base64.b64decode(image_data.img_bytes)
|
|
results = ocr_results(image_bytes)
|
|
|
|
# Explicitly delete unused variables and run garbage collector
|
|
del image_bytes
|
|
gc.collect()
|
|
|
|
return {"results": results}
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import uvicorn
|
|
|
|
uvicorn.run(app, host="127.0.0.1", port=8000)
|