表计数据自动化添加分割刻度，ocr画框

Datainterpolation_and_Ocrframe.py

表计自动化标注和ocr画框

备份记录

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
 * Copyright (c) 2023 XiangYunShuZhi Corporation.
 * @brief      数据处理，可实现自动化添加分割刻度，ocr画框
 * @author     lll
 * @date       2023/12/08
 * @history    修改不规范代码,增加代码注释
"""
import re
import os
import json


def insert_point(points, points_num: int, scale: int, n, m_list, mul_w_list, mul_h_list, route):
    """
    根据自定义的点个数，刻度标准，延长线长度来在图像中插入间隔刻度

    :param points: 根据json读取的初始标注点位信息，无需更改
    :param points_num: 初始标注点位的个数，根据实际标注情况调整，一般为12点或14点
    :param scale: 刻度标准，一般设置为10刻度标准或5刻度标准
    :param n:  延长线是初始线段的n倍，可根据实际情况调整，一般设置0.3、0.5等
    :param m_list:  标注框中心点到初始上点组成的线段是初始线段的m倍，一般是延长线的2倍
    :param mul_w_list:  矩形数字框的宽度是初始线段的w倍
    :param mul_h_list:  矩形数字框的高度是初始线段的h倍
    :param route: 遍历文件夹下每个json的路径
    l1: 延长线长度
    l2: 初始标注点位对应的两个上下点组成的线段长度
    l3: 初始标注点位对应的上点到标注框中心点组成的线段长度

    :return: 初始点位的延长线段l3定位出的ocr数字框点位信息
    """
    up_points = points[int(points_num / 2):]
    down_points = points[:int(points_num / 2)]
    num_up_points = points[int(points_num / 2):]
    num_down_points = points[:int(points_num / 2)]
    rec_left_up = []
    rec_right_up = []
    rec_right_down = []
    rec_left_down = []
    rec_points_list = []
    print('len(up_points):', len(up_points))
    print('len(down_points):', len(down_points))

    # 设置分割间隔
    for i in range(len(up_points) - 1):
        up_gap_x = (up_points[i + 1][0] - up_points[i][0]) / scale
        up_gap_y = (up_points[i + 1][1] - up_points[i][1]) / scale
        down_gap_x = (down_points[i + 1][0] - down_points[i][0]) / scale
        down_gap_y = (down_points[i + 1][1] - down_points[i][1]) / scale

        for j in range(scale - 1):
            up_point = [up_points[i][0] + (j + 1) * up_gap_x, up_points[i][1] + (j + 1) * up_gap_y]
            up_points.append(up_point)

            down_point = [down_points[i][0] + (j + 1) * down_gap_x, down_points[i][1] + (j + 1) * down_gap_y]
            down_points.append(down_point)

    # 上下点按顺序排序
    up_points.sort(key=lambda x: x[0])
    down_points.sort(key=lambda x: x[0])
    # 新增num_up_points用来添加识别的数字框
    num_up_points.sort(key=lambda x: x[0])
    num_down_points.sort(key=lambda x: x[0])

    # 计算l1延长线和l2初始线段的长度
    for i in range(len(up_points)):
        if i == 0 or i % scale == 0:
            #   x0 - x1
            compare_x = (down_points[i][0] - up_points[i][0])
            # 求l2
            y0 = down_points[i][1]
            x0 = down_points[i][0]
            y1 = up_points[i][1]
            x1 = up_points[i][0]
            l2pf = (y0 - y1) ** 2 + (x0 - x1) ** 2    # 计算l2的平方
            l2 = l2pf ** 0.5    # 开根号，计算l2的长度
            l1 = n * l2    # 设置l1为n倍l2
            # 求y2 ,               l1,延长线的长度
            y2 = y0 - (((y0 - y1) * (l1 + l2)) / l2)
            if compare_x >= 0:
                up_point = [x0 - ((x0 - x1) * (l1 + l2)) / l2, y2]
                up_points[i] = up_point
            if compare_x < 0:
                up_point = [(((x1 - x0) * (l1 + l2)) / l2) + x0, y2]
                up_points[i] = up_point

    # 增加l3延长线作为数字框的中心点
    for num, m, mul_w, mul_h in zip(range(len(num_up_points)), m_list, mul_w_list, mul_h_list):
        #   x0 - x1
        compare_x = (num_down_points[num][0] - num_up_points[num][0])
        # 求l2
        y0 = num_down_points[num][1]
        x0 = num_down_points[num][0]
        y1 = num_up_points[num][1]
        x1 = num_up_points[num][0]
        l2pf = (y0 - y1) ** 2 + (x0 - x1) ** 2
        l2 = l2pf ** 0.5
        l3 = m * l2
        w = mul_w * l2
        h = mul_h * l2

        # 求y3 ,               l3：到数字框延长线中点的长度
        y3 = y0 - (((y0 - y1) * (l3 + l2)) / l2)
        if compare_x >= 0:
            num_up_point = [x0 - ((x0 - x1) * (l3 + l2)) / l2, y3]
            num_up_points[num] = num_up_point

        if compare_x < 0:
            num_up_point = [(((x1 - x0) * (l3 + l2)) / l2) + x0, y3]
            num_up_points[num] = num_up_point

        rec_left_up_x = num_up_points[num][0] - (0.5 * w)
        rec_left_up_y = num_up_points[num][1] - (0.5 * h)
        rec_right_up_x = num_up_points[num][0] + (0.5 * w)
        rec_right_up_y = num_up_points[num][1] - (0.5 * h)
        rec_right_down_x = num_up_points[num][0] + (0.5 * w)
        rec_right_down_y = num_up_points[num][1] + (0.5 * h)
        rec_left_down_x = num_up_points[num][0] - (0.5 * w)
        rec_left_down_y = num_up_points[num][1] + (0.5 * h)
        rec_left_up.append([rec_left_up_x, rec_left_up_y])
        rec_right_up.append([rec_right_up_x, rec_right_up_y])
        rec_left_down.append([rec_left_down_x, rec_left_down_y])
        rec_right_down.append([rec_right_down_x, rec_right_down_y])

    # 将生成的上下点信息写到json里
    for i in range(len(up_points)):
        tmp = {
            "label": "line",
            "points": [
                up_points[i],
                down_points[i]
            ],
            "group_id": None,
            "shape_type": "line",
            "flags": {}
        }
        data['shapes'].append(tmp)

    with open(f'{route}', 'w', encoding='utf-8') as f:
        json.dump(data, f, ensure_ascii=False, indent=2)

    # 初始点位的延长线段l3定位出的ocr数字框
    for i in range(len(num_up_points)):
        num_tmp = [{
            "points": [
                rec_left_up[i],
                rec_right_up[i],
                rec_right_down[i],
                rec_left_down[i]
            ]}
        ]
        rec_points = num_tmp[0]['points']
        rec_points_list.append(rec_points)

    return rec_points_list


def list_json_files(directory: str):
    """
    遍历指定目录下的所有JSON文件并打印它们的文件名

    :param:  directory: 指定要遍历的目录
    :return: dir_json:  文件夹下所有的json文件组成的列表
    """
    dir_json = []
    for root, dirs, files in os.walk(directory):
        for file in files:
            if file.endswith('.json'):
                dir_json.append(os.path.join(root, file))

    return dir_json


def list_jpg_files(directory: str):
    """
    遍历指定目录下的所有jpg文件并打印它们的文件名

    :param:  directory: 指定要遍历的目录
    :return: dir_jpg:  文件夹下所有的jpg文件组成的列表
    """
    dir_jpg = []
    jpg_name = []
    for root, dirs, files in os.walk(directory):
        for file in files:
            if file.endswith('.jpg'):
                dir_jpg.append(os.path.join(root, file))
                jpg_name.append(file)

    return dir_jpg


def change_quotes(text):
    """
    查找文本中的单引号，
    将单引号变成双引号。
    :param text: 读取的文本信息
    :return: 替换成双引号的文本信息
    """

    single_quotes = re.findall(r'(?<!`)(\'[^\']*\')', text)

    for i, single_quote in enumerate(single_quotes):
        double_quote = single_quote.replace("'", '"')
        text = text.replace(single_quote, double_quote)
    return text


def read_file(file_path):
    """
    读取文件内容
    :param file_path: 本脚本同目录下生成的Label.txt文件路径
    :return: 读取的文本信息
    """
    with open(file_path, 'r', encoding='utf-8') as f:
        text = f.read()
    return text


def write_file(file_path, text):
    """
    写入文件内容
    :param file_path: 本脚本同目录下生成的Label.txt文件路径
    :param text: 要写入的文本信息
    """
    with open(file_path, 'w', encoding='utf-8') as f:
        f.write(text)


def gen_points_info(points, num):
    """
    # 获取Label.txt文件要中写入的数字框信息

    :param points: 初始点位的延长线段l3定位出的ocr数字框
    :param num:    图片上的数字刻度
    :return:       要写入文本的数字框信息
    """
    tmp = list()
    for i, v in enumerate(points):
        # if i == 0 or i % 2 == 0:  (本条判断用来指定哪些数字框是需要写入的)
        tmp.append(
            {
                "transcription": num[i],
                "points": v,
                "difficult": False
            }
        )
    return tmp


def run_ocr(scale_num, dir_name):
    """
    执行ocr画数字框
    :param scale_num: 图片上的数字刻度
    :param dir_name:  Label.txt文件中图片名称的路径文件夹前缀
    """
    for q in range(len(get_rec_points_list)):
        for w in range(len(get_rec_points_list[q])):
            for e in range(len(get_rec_points_list[q][w])):
                for r in range(len(get_rec_points_list[q][w][e])):
                    # 对每个数字进行四舍五入取整操作
                    get_rec_points_list[q][w][e][r] = round(get_rec_points_list[q][w][e][r])

    get_dirjpg = list_jpg_files(directory)
    # 将图片名称写入txt文件
    jpg_file_name_list = []
    with open('fileState.txt', 'w') as f2:
        for jpg_file in get_dirjpg:
            f2.write(os.path.join(directory, jpg_file) + '\t' + '1' + '\n')
            jpg_file_name = jpg_file.split('\\')[-1]
            jpg_file_name_list.append(jpg_file_name)
    print("fileState.txt : 已生成")

    # scale_num = ["0", "20", "40", "60", "80", "100", "120", "140", "160"]
    with open("Label.txt", "w") as f:
        for jpg, i_rec_points in enumerate(get_rec_points_list):
            txt1 = jpg_file_name_list[jpg]
            txt2 = gen_points_info(i_rec_points, scale_num)
            txt3 = json.dumps(txt2)
            f.write(f"{dir_name}/" + txt1 + "\t" + txt3 + "\n")

    print("Label.txt : 已生成")

    file_path = "Label.txt"
    text = read_file(file_path)
    new_text = change_quotes(text)

    write_file(file_path, new_text)


if __name__ == '__main__':
    # 运行程序需要更改的步骤1，指定要遍历的目录
    directory = r'D:\pythonProject\1206-differentauto\1206\20231206-yun'

    # 调用函数并打印JSON文件的文件名
    get_dirjson = list_json_files(directory)
    j = 0
    get_rec_points_list = []
    print('1、执行刻度分割-----------------------')
    for i in get_dirjson:
        with open(f'{i}', 'r', encoding='utf-8') as f1:
            data = json.load(f1)
        # 获取"points"字段
        org_points = data['shapes'][0]['points']
        # 清掉json里初始'shapes'的点位
        data['shapes'].remove(data['shapes'][0])

        # points_num    是初始标注点个数
        # scale         是刻度标注，类似中间插值个数+1
        # n=1.2         延长线1的长度是初始上下点位组成的线段长度的倍数
        # m_list        延长线2（到ocr数字框的中心点）的长度是初始上下点位组成的线段长度的倍数
        # mul_w_list    ocr数字框的宽度是初始上下点位组成的线段长度的倍数
        # mul_h         ocr数字框的高度是初始上下点位组成的线段长度的倍数
        # 不需要更改的：
        # ori_points 原始标注文件的点位信息
        # route  是遍历的文件夹中json的路径
        # 运行程序需要更改的步骤2：
        get_rec_point = insert_point(org_points, points_num=int(18), scale=10, n=1,
                                     m_list=[-2, -2.1, -2.2, -2.1, -2, -2.1, -2.5, -2.4, -2.4],
                                     mul_w_list=[1.3, 2, 2, 2, 2, 3, 3, 3, 3], mul_h_list=[2, 2, 2, 2, 2, 2, 2, 2, 2], route=i)
        get_rec_points_list.append(get_rec_point)
        print(f'{i} 分割插入完成')
        print('~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~')

        j += 1

    print(f'该文件夹下共{j}个json文件已插入更新')
    # 是否执行ocr
    print('2、执行orc画框-----------------------')
    run_ocr(scale_num=["0", "20", "40", "60", "80", "100", "120", "140", "160"], dir_name='20231206-yun')