pdf_title_image/ocr_api.py

# ocr外部接口
import os
from typing import List

from requests import post


class OcrAgent:
    def __init__(self, url):
        self.url = url

    def get_content(self, image_path):
        try:
            with open(image_path, 'rb') as image_file:
                files = {"file": ("image.jpg", image_file, "image/jpeg")}
                # files = {"file": ("image.png", image_file, "image/png")}
                response = post(self.url, files=files)
            return response.json()
        except ValueError:
            raise ValueError(f"传入图像{image_path}已损坏")


def find_current_row(ocr_result: List[dict], top: int, bottom: int, float_range: int = 5):
    results = []
    assert float_range >= 0
    top += float_range
    bottom -= float_range
    for ret in ocr_result:
        ct = ret['rect']['top']
        cb = ret['rect']['top'] - ret['rect']['height']
        if top >= ct > cb >= bottom:
            results.append(ret)
    return results




if __name__ == '__main__':
    agent = OcrAgent("http://120.48.103.13:18000/ctr_ocr")
    res = agent.get_content(
        os.path.join('/home/zzh/ocr/pdf/南方电网数字电网研究院有限公司/images', 'image_page_1131_0.png'))
    pass