yl-ocr-layout/create_md.py

"""
function: 生成md文件
"""
from mdutils.mdutils import MdUtils
import os
import cv2
import numpy as np
import requests

d_map = {0: 'code', 1: 'logo', 2: 'style', 3: 'table', 4: 'text', 5: 'title'}
images_path = "./train/images/"
labels_path = "./train/labels/"
cut_ori_img_path = "./train/cut_ori_imgs/"
cut_predict_img_path = "./train/cut_predict_imgs/"
md_cut_ori_img_path = "../train/cut_ori_imgs/"
md_cut_predict_img_path = "../train/cut_predict_imgs/"
md_path = "mdfiles/"


# 标签是经过归一化的，需要变回来
# xywh格式 ---> box四个顶点的坐标
def xywh2lrbt(img_w, img_h, box):
    c, x, y, w, h = int(box[0]), float(box[1]), float(box[2]), float(box[3]), float(box[4])
    x = x * img_w  # 中心坐标x
    w = w * img_w  # box的宽
    y = y * img_h  # 中心坐标y
    h = h * img_h  # box的高

    lt_x, lt_y = x - w / 2, y - h / 2  # left_top_x, left_top_y
    lb_x, lb_y = x - w / 2, y + h / 2  # left_bottom_x, left_bottom_y
    rb_x, rb_y = x + w / 2, y + h / 2  # right_bottom_x, right_bottom_y
    rt_x, rt_y = x + w / 2, y - h / 2  # right_top_x, right_bottom_y

    lrbt = [[lt_x, lt_y], [lb_x, lb_y], [rb_x, rb_y], [rt_x, rt_y]]

    return lrbt, c


def IOU(rect1, rect2):
    xmin1, ymin1, xmax1, ymax1 = rect1
    xmin2, ymin2, xmax2, ymax2 = rect2
    s1 = (xmax1 - xmin1) * (ymax1 - ymin1)
    s2 = (xmax2 - xmin2) * (ymax2 - ymin2)

    sum_area = s1 + s2

    left = max(xmin2, xmin1)
    right = min(xmax2, xmax1)
    top = max(ymin2, ymin1)
    bottom = min(ymax2, ymax1)

    if left >= right or top >= bottom:
        return 0

    intersection = (right - left) * (bottom - top)
    return intersection / (sum_area - intersection) * 1.0


def send_requests(img_path):
    b_img_list = []
    b_img_list.append(('file_list', ('image.jpeg', open(img_path, 'rb'), 'image/jpeg')))
    payload = {'model_name': 'ocr-layout', 'img_size': 1920, 'download_image': False}
    response = requests.request("POST", "http://192.168.199.249:4869/detect", headers={}, data=payload,
                                files=b_img_list)
    result = eval(str(response.text))
    result = eval(str(result['result']))
    return result


if __name__ == "__main__":
    list_images = os.listdir(images_path)
    list_labels = os.listdir(labels_path)
    all_acc = 0
    all_label_len = 0

    # 遍历每张图片
    for m in range(len(list_images)):
        each_acc_map = {0: 0, 1: 0, 2: 0, 3: 0, 4: 0, 5: 0}  # 每张图片对应的各自预测正确数量
        each_acc = 0  # 每张图片总共预测正确的数量

        image = list_images[m]
        img_name = image.split(".")[0]
        # if image != "125-------------_jpg.rf.9a0d958c405ffb596b4022e13f0686b7.jpg":
        #     continue
        for n in range(len(list_labels)):
            label = list_labels[n]
            label_name = label.split(".")[0]
            if img_name != label_name:
                continue
            if img_name == label_name:
                print(image)
                print(label)
                print("===========")
                # 每张图都创建一个md文件
                mdFile = MdUtils(file_name=md_path+'yq_0819_' + str(img_name), title='yq_0819_' + str(img_name) + '_iou>=0.5')
                list_item = ["ori_cla", "predict_cla", "cut_ori_img", "cut_predict_img", "isTrue"]

                result = send_requests(images_path + image)

                img = cv2.imread(images_path + image)
                img_h, img_w, _ = img.shape
                # 获取原图上的切割结果
                arr_label = []  # labels的边界框，存储四个顶点坐标
                ori_clas = []  # label的真实类别
                with open(labels_path + label, mode="r", encoding="utf-8") as f:
                    lines = f.readlines()
                for line in lines:
                    line = line.split(" ")
                    lrbt, c = xywh2lrbt(img_w, img_h, line)
                    arr_label.append(lrbt)
                    ori_clas.append(c)

                all_label_len += len(lines)  # 所有标签的数量，为了后面计算所有图片的准确度

                num = 0
                for i in range(len(arr_label)):
                    lt_x, lt_y = int(arr_label[i][0][0]), int(arr_label[i][0][1])
                    rb_x, rb_y = int(arr_label[i][2][0]), int(arr_label[i][2][1])
                    rect1 = [lt_x, lt_y, rb_x, rb_y]

                    # 切割真实标签，并保存到本地
                    cut_ori_img = img[lt_y:rb_y, lt_x:rb_x, :]
                    save_cut_ori_img = cut_ori_img_path + str(img_name) + "_" + str(num) + ".jpg"
                    cv2.imwrite(save_cut_ori_img, cut_ori_img)
                    save_cut_ori_img = mdFile.new_inline_image(text='', path=md_cut_ori_img_path+ str(img_name) + "_" + str(num) + ".jpg")
                    ori_cla = int(ori_clas[i])
                    print("此时的ori_cla:", d_map[ori_cla])
                    num += 1

                    for j in range(len(result[0])):
                        res = []
                        predict_cla = int(result[0][j]['class'])
                        print("此时的predict_cla:", d_map[predict_cla])
                        if ori_cla == predict_cla or (ori_cla in [2, 3, 4] and predict_cla in [2, 3, 4]):
                            rect2 = result[0][j]['bbox']
                            iou = IOU(rect1, rect2)
                            if iou >= 0.5:
                                # 切割预测结果
                                cut_predict_img = img[rect2[1]:rect2[3], rect2[0]:rect2[2], :]
                                save_cut_predict_img = cut_predict_img_path + str(img_name) + "_" + str(
                                    j) + ".jpg"
                                cv2.imwrite(save_cut_predict_img, cut_predict_img)
                                save_cut_predict_img = mdFile.new_inline_image(text='', path=md_cut_predict_img_path+ str(img_name) + "_" + str(j) + ".jpg")
                                isTrue = True
                                list_item.extend(
                                    [d_map[ori_cla], d_map[predict_cla], save_cut_ori_img, save_cut_predict_img,
                                     isTrue])
                                each_acc_map[ori_cla] += 1
                                all_acc += 1
                                each_acc += 1
                                break
                        if j == len(result[0]) - 1:
                            isTrue = False
                            list_item.extend([d_map[ori_cla], " ", save_cut_ori_img, " ", isTrue])
                            list_item.extend(res)

                # 每一张图片的总准确度
                each_acc = round(each_acc / len(ori_clas), 3)
                # 每一类的准确度
                code = round(each_acc_map[0] / ori_clas.count(0), 3) if ori_clas.count(0) else "这张图没有这个标签"
                logo = round(each_acc_map[1] / ori_clas.count(1), 3) if ori_clas.count(1) else "这张图没有这个标签"
                style = round(each_acc_map[2] / ori_clas.count(2), 3) if ori_clas.count(2) else "这张图没有这个标签"
                table = round(each_acc_map[3] / ori_clas.count(3), 3) if ori_clas.count(3) else "这张图没有这个标签"
                text = round(each_acc_map[4] / ori_clas.count(4), 3) if ori_clas.count(4) else "这张图没有这个标签"
                title = round(each_acc_map[5] / ori_clas.count(5), 3) if ori_clas.count(5) else "这张图没有这个标签"

                mdFile.new_line(str(image) + "—总准确率:" + str(each_acc))
                mdFile.new_line(str(d_map[0]) + "—准确率:" + str(code))
                mdFile.new_line(str(d_map[1]) + "—准确率:" + str(logo))
                mdFile.new_line(str(d_map[2]) + "—准确率:" + str(style))
                mdFile.new_line(str(d_map[3]) + "—准确率:" + str(table))
                mdFile.new_line(str(d_map[4]) + "—准确率:" + str(text))
                mdFile.new_line(str(d_map[5]) + "—准确率:" + str(title))

                break

        mdFile.new_line()
        mdFile.new_table(columns=5, rows=len(list_item) // 5, text=list_item, text_align='center')
        mdFile.create_md_file()

    print(all_acc, all_label_len)
    print("所有图片的预测准确度：", all_acc / all_label_len)