exam-segment-django/segment/ocr/penguin_ocr.py

# @Author  : lightXu
# @File    : penguin_ocr.py
# @Time    : 2019/6/11 0011 下午 17:52
import base64
import hashlib
import random
import string
import time
from urllib.parse import urlencode

import cv2
import requests
from requests.adapters import HTTPAdapter
from urllib3.util.retry import Retry

from segment.image_operation.utils import resize_by_percent, write_single_img

APP_KEY = 'R2iPkd5J2056YFRw'
APP_ID = '2117302084'


def opecv2base64(img):
    image = cv2.imencode('.jpg', img)[1]
    base64_data = str(base64.b64encode(image))[2:-1]
    return base64_data


def get_base64_size(base64_str):
    length = len(base64_str)
    size = float(length - 2 * length/8)  # byte
    return size


def opecv2base64_stand(raw_image, mem_size, default_size=1):  # 小于1M
    default_size = default_size * 1000 * 1000
    m_ratio = mem_size/default_size
    if m_ratio > 1.0:
        y, x = raw_image.shape[0], raw_image.shape[1]
        s_ratio = max(y, x) / 1200
        ratio = max(m_ratio, s_ratio)
        image_resize = resize_by_percent(raw_image, 1/ratio)
        # img_gray = cv2.cvtColor(raw_image, cv2.COLOR_RGB2GRAY)
        write_single_img(image_resize, r'C:\Users\Administrator\Desktop\p\01_r.jpg')
        return opecv2base64(image_resize)
    else:
        return opecv2base64(raw_image)


def _get_sign(params, app_key):
    sort_dict = sorted(params.items(), key=lambda item: item[0], reverse=False)
    sort_dict.append(('app_key', app_key))
    rawtext = urlencode(sort_dict).encode()
    sha = hashlib.md5()
    sha.update(rawtext)
    md5text = sha.hexdigest().upper()

    return md5text


def please_retry(response, url, data, headers):
    status_code = response.status_code
    if status_code == 200:
        resp = response.json()
        if 'ok' != resp.get('msg'):

            try_iter = 0
            while try_iter < 3:
                response = requests.post(url, data=data, headers=headers, timeout=15)

            print(resp)


def get_ocr_english_text_raw_format(img, size):
    url = 'https://api.ai.qq.com/fcgi-bin/ocr/ocr_generalocr'
    headers = {'Content-Type': 'application/x-www-form-urlencoded'}
    image_base64 = opecv2base64_stand(img, size)  # 得到 base64 编码的数据
    nonce_str = ''.join(random.sample(string.ascii_letters + string.digits, 10))
    data = {
        'app_id': APP_ID,
        'image': image_base64,
        'time_stamp': str(int(time.time())),
        'nonce_str': nonce_str,
    }

    sign = _get_sign(data, APP_KEY)
    data['sign'] = sign

    s = requests.Session()
    retries = Retry(total=3,
                    backoff_factor=0.1,
                    status_forcelist=[500, 502, 503, 504])
    s.mount('https://', HTTPAdapter(max_retries=retries))

    # response = requests.post(url, data=data, headers=headers)
    response = s.post(url, data=data, headers=headers, timeout=15)

    final_response = ''
    for _ in range(0, 3):
        status_code = response.status_code
        if status_code == 200 and 'ok' == response.json().get('msg'):
            final_response = response
            # print('ok')
            break
        else:
            response = s.post(url, data=data, headers=headers, timeout=15)
            # print('retry')

    if final_response:
        status_code = final_response.status_code
        if status_code == 200:
            resp = response.json()
            if 'ok' != resp.get('msg'):
                # print(resp)
                raise Exception("ocr error {}: {}!".format(resp.get('ret'), resp.get('msg')))
            # print(resp)
        else:
            raise ValueError('ocr failed, response[{}]'.format(status_code))
    else:
        raise ValueError('ocr failed, retried three times while no response')

    return resp


def ocr_format(resp):
    item_list = resp['data']['item_list']
    words_str_list = []
    for item_index, item in enumerate(item_list):
        words_list = item['words']
        words_str = ''
        for char_index, char_dict in enumerate(words_list):
            char = char_dict['character']
            if char == '':
                char = ' '
            words_str = words_str + char
        words_str_list.append(words_str.lstrip())

    return words_str_list


def get_ocr_english_text(image, size):
    resp = get_ocr_english_text_raw_format(image, size)
    words_list = ocr_format(resp)
    return words_list