hr-ocr-idcard/generate_test.py

import sys
import re
import json
from itertools import chain
from pathlib import Path
import cv2

# 配置


img_type = 0
data = '{"image": img_str, "image_type": image_type}'
test_path = 'testing/id727_0_test.py'
images_path = Path('images/test_img/7.27/0')
root = Path(__file__).parent
img_paths = chain(*[Path(root / images_path).rglob(f'*.{ext}') for ext in ['jpg']])
dict_orientation = {0: None, 1: cv2.ROTATE_90_CLOCKWISE, 2: cv2.ROTATE_180, 3: cv2.ROTATE_90_COUNTERCLOCKWISE}
head = f"""
import unittest
import base64
from dataclasses import dataclass
from pathlib import Path
import cv2
import requests

url = 'http://192.168.199.249:18080'


def send_request(image_path, image_type, rotate=None):
    img = cv2.imread(str(image_path))
    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    if rotate is not None:
        img = cv2.rotate(img, rotate)
    _, im_arr = cv2.imencode('.jpg', img)
    img_str = base64.b64encode(im_arr).decode('utf-8')
    r = requests.post(url + '/ocr_system/idcard', json={data})
    print(r.json())
    return r.json()

@dataclass
class ResultItem:
    status: str
    orientation: int
    name: str
    id: str
    ethnicity: str
    gender: str
    birthday: str
    address: str
    address_province: str
    address_city: str
    address_region: str
    address_detail: str
    expire_date:str
    """
masterclass = f"""
class TestSchoolCertOcr(unittest.TestCase):
    def _helper(self, image_path, item: ResultItem,  rotate=None,image_type={img_type}):
        root = Path(__file__).parent
        image_path = str(root / image_path)
        r = send_request(image_path, image_type, rotate)
        self.assertEqual(item, ResultItem(status=r['status'],
                                          orientation=r['result']['orientation'],
                                          name=r['result']['name']['text'],
                                          id=r['result']['id']['text'],
                                          ethnicity=r['result']['ethnicity']['text'],
                                          gender=r['result']['gender']['text'],
                                          birthday=r['result']['birthday']['text'],
                                          address=r['result']['address']['text'],
                                          address_province=r['result']['address_province']['text'],
                                          address_city=r['result']['address_city']['text'],
                                          address_region=r['result']['address_region']['text'],
                                          address_detail=r['result']['address_detail']['text'],
                                          expire_date=r['result']['expire_date']['text']))
"""

with open(root / test_path, 'w', encoding='utf-8') as f:
    f.write(head)
    f.write(masterclass)
    for path in img_paths:
        pic_name = path.stem.split("_img")[0]
        pic_path = str(path.parent).split("images")[-1] + '/' + path.name
        pic_json_name = path.stem + '.json'
        pic_json_path = path.parent / pic_json_name

        with pic_json_path.open('r') as json_f:
            json_input = json.load(json_f)
            del [json_input['card_type']]
            statsu_input = {'status': '000'}
            json_input = dict(statsu_input, **json_input)
            testfunc = f"""
    def case_{pic_name}_0(self, orientation=0):
        dict_orientation = {dict_orientation}
        image_path = '../images{pic_path}'
        json_input = {json_input}
        json_input['orientation'] = orientation
        self._helper(image_path, ResultItem(**json_input), dict_orientation[orientation])

    def test_case_{pic_name}_0(self):
        self.case_{pic_name}_0(0)    

    def test_case_{pic_name}_1(self):
        self.case_{pic_name}_0(1)

    def test_case_{pic_name}_2(self):
        self.case_{pic_name}_0(2)

    def test_case_{pic_name}_3(self):
        self.case_{pic_name}_0(3)

            """
        f.write(testfunc)

print('ok')