Jan 10, 2019

python 利用百度 AI 实现 OCR

前两天了解了一下百度 AI，他们提供了免费的图片中文字识别接口，每天可以提供 50000 次免费的通用文字识别，想起了之前帮水木做一个网站时，资料中的问题大部分以图片的形式提供，所以当时就有这个图片 OCR 的需求，本文就说一下如何用 python 借用百度 AI 实现图片的 OCR。

实现分析

阅读了百度提供的文章只要10分钟 快速掌握文字识别，基本就可以整理出实现的步骤了：

  1. 创建百度应用，获取 api_key 和 secret_key
  2. 利用 api_key 和 secret_key 获取 access token
  3. 图片转 BASE64
  4. 利用图片BASE64码和 access token 构建 POST
  5. 响应结果中解析结果

其中第一步需要在百度 AI 控制台实现，有了 api_key 和 secret_key 剩下的四步就可以用 python 实现了，最终 python 程序设计思路：

  1. 指定配置文件和图片路径
  2. 读取配置文件中的 api_key 和 secret_key
  3. 判断配置文件是否有 access_token
    • 如果没有：利用 api_key 和 secret_key 获取 access_key，并将 access_token 写入配置文件
    • 如果有：读取配置中的 access_key
  4. 将指定图片转换为 BASE64 码
  5. 根据 access_token 和 图片的 BASE64 码 构建 POST，发送请求，解析返回的结果
  6. 打印 OCR 结果

代码实现

textGO 类

根据上面的思路，封装了 textGO 类：

class TextGO:
    """ 定义 ocr 类 """

    def __init__(self, config_file):
        config = self.read_config(config_file)
        if not (config.get("api_key") and config.get("secret_key")):
            print("请指定一个合理的配置文件，配置文件中要包含 api_key 和 secret_key!")
            self.access_token = ""
            return
        access_token = config.get("access_token")
        if access_token is None or access_token == "":
            access_token = self.get_access_token(
                config.get("api_key"), config.get("secret_key"))
            config["access_token"] = access_token
            self.write_config(config_file, config)
        self.access_token = config.get("access_token")

    @staticmethod
    def read_config(config_file) -> dict:
        """ 从 json 文件中读取配置 """
        if not os.path.exists(config_file):
            print(f"指定的配置文件 {config_file} 不存在!")
            return {}
        with open(config_file, "r") as load_file:
            return json.load(load_file)
        return {}

    @staticmethod
    def write_config(config_file, config):
        """ 将配置写入文件 """
        with open(config_file, "w") as f_obj:
            json.dump(config, f_obj)

    @staticmethod
    def img_to_base64(img):
        """ convert img to base64 data """
        with open(img, "rb") as img_f:
            file_data = img_f.read()
            b64_data = b64encode(file_data)
            b64_str = str(b64_data, "utf-8")
        return b64_str

    @staticmethod
    def get_access_token(api_key, secret_key):
        """ 获取百度 AI 的文字识别 access_token """
        print("获取百度 AI 文字识别的 Access Token ...")
        url = "https://aip.baidubce.com/oauth/2.0/token?grant_type=client_credentials&client_id=%s&client_secret=%s"
        response: requests.models.Response = requests.get(
            url % (api_key, secret_key))
        body_str = str(response.content, "utf-8")
        body_dict = json.loads(body_str)
        return body_dict.get("access_token")

    def ocr(self, img):
        """ 利用百度 AI 识别文字 """
        if self.access_token == "":
            print("没有 access_token")
            return ""
        url = f"https://aip.baidubce.com/rest/2.0/ocr/v1/accurate?access_token={self.access_token}"
        headers = {"Content-Type": "application/x-www-form-urlencoded"}
        data = {
            "image_type": "BASE64",
            "image": self.img_to_base64(img),
            "group_id": "textGO_test",
            "user_id": "5km"
        }
        response = requests.post(url=url, data=data, headers=headers)
        if response.status_code == 200:
            body_str = str(response.content, "utf-8")
            body_dict = json.loads(body_str)
            words = ""
            for line in body_dict.get("words_result"):
                words += line.get("words") + "\n"
            return words
        return ""
  1. 图片转码使用了内建库 base64
  2. 配置的读取和写入使用了库 json
  3. 发送 GET 和 POST 请求使用库 requests

解析命令参数

为了更好的实现命令形式，这里使用了库 argparse，封装了函数 get_args：

def get_args() -> argparse.Namespace:
    """ 获取命令参数 """
    parser = argparse.ArgumentParser(description="OCR 识别图片中的文字")
    parser.add_argument("img", help="指定要识别的图片")
    parser.add_argument(
        "-c", "--config", default="./config.json", help="指定配置文件")
    return parser.parse_args()

功能使用

定义 main 函数，实现整个过程：

def main():
    """ 主函数 """
    args = get_args()
    textGO = TextGO(args.config)
    words = textGO.ocr(args.img)
    print(words)

在使用前还需要按照只要10分钟 快速掌握文字识别中的操作获取自己的 api_key 和 secret_key ，将他们以下面的形式写入到一个文件，我这里写入文件是 config.json：

{"api_key": "", "secret_key": ""}

将自己应用的 api_key 和 secret_key 对应填到上面 json 内容的对应位置！

然后调用主函数：

if __name__ == "__main__":
    main()

最终 textGO.py 的文件代码就完成了。

测试功能

python3 textGO.py -h

通过上述命令查看帮助信息：

usage: textGO.py [-h] [-c CONFIG] img

OCR 识别图片中的文字

positional arguments:
  img                   指定要识别的图片

optional arguments:
  -h, --help            show this help message and exit
  -c CONFIG, --config CONFIG
                        指定配置文件

准备一张图片，假如：

  • 图片路径是 ~/Desktop/demo.png：

  • 配置文件路径为 ~/Desktop/config

那么可以使用如下命令并得到结果:

python3 textGO.py -c ~/Desktop/config.json ~/Desktop/demo.png
获取百度 AI 文字识别的 Access Token ...
 macOS Mojave
版本10.14.2
 MacBook Pro(15-inch, 2016)
处理器2.6 GHz Intel Core i7
内存16GB2133 MHZ LPDDR3
图形卡 Radeon Pro4502048MB
 Intel HD Graphics 530 1536 MB
序列号co2 SXCQWGTFL

代码托管

代码已托管至 github：tools-with-script/textGO/

访问量