hr-ocr-regbook/generate_test.py

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

# 配置


img_type = 1
data = '{"image": img_str, "image_type": image_type}'
test_path = 'testing/clear_1_test.py'
images_path = Path('images/clean/1')
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:38811'


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/regbook', json={data})
    print(r.json())
    return r.json()

@dataclass
class ResultItem:
    status: str
    orientation: int
    type: str
    address: str
    address_province: str
    address_city: str
    address_region: str
    address_detail: str
    name: str
    id: str
    gender: str
    birthplace: str
    birthplace_province: str
    birthplace_city: str
    birthplace_region: str
    native_place: str
    native_place_province: str
    native_place_city: str
    native_place_region: str
    blood_type: str
    religion: str
    """
masterclass = f"""
class TestRegBookOcr(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'],
                                          type=r['result']['type']['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'],
                                          name=r['result']['name']['text'],
                                          id=r['result']['id']['text'],
                                          gender=r['result']['gender']['text'],
                                          birthplace=r['result']['birthplace']['text'],
                                          birthplace_province=r['result']['birthplace_province']['text'],
                                          birthplace_city=r['result']['birthplace_city']['text'],
                                          birthplace_region=r['result']['birthplace_region']['text'],
                                          native_place=r['result']['native_place']['text'],
                                          native_place_province=r['result']['native_place_province']['text'],
                                          native_place_city=r['result']['native_place_city']['text'],
                                          native_place_region=r['result']['native_place_region']['text'],
                                          blood_type=r['result']['blood_type']['text'],
                                          religion=r['result']['religion']['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['img_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')