yl-ocr-layout/writeMD.py

from mdutils.mdutils import MdUtils
import os
import cv2
import numpy as np
import requests


# url = "http://192.168.199.249:4869"

# 标签是经过归一化的，需要变回来
# xywh格式 ---> box四个顶点的坐标
def xywh2lrbt(img_w, img_h, box):
    c, x, y, w, h = float(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():
    return result

img_path = "./images/test1.jpg"
b_img_list = []
b_img_list.append(('file_list', ('image.jpeg', open(img_path, 'rb'), 'image/jpeg')))
payload = {'model_name': 'ocr-layout', 'img_size': 800, 'download_image': False}
response = requests.request("POST", "http://192.168.199.107:18089/detect", headers={}, data=payload, files=b_img_list)
result = eval(str(response.text))
result = eval(str(result['result']))
# print(len(result))

mdFile = MdUtils(file_name='yolov5_yili_0818_1', title='Markdown File Example')
mdFile.new_header(level=1, title="yolov5预测结果")
list_item = ["ori_cla", "predict_cla", "cut_ori_img", "cut_predict_img", "isTrue"]
# list_item.extend(["ori_cla", "predict_cla", "cut_ori_img", "cut_predict_img", "isTrue"])
# for i in range(len(result[0])):
#     # print(stencil_list[0][i])
#
#     predict_cla = result[0][i]["class"]

# 获取真实的标签
label = "./images/test1.txt"
with open(label, mode="r", encoding="utf-8") as f:
    lines = f.readlines()

img = cv2.imread(img_path)
img_h, img_w, _ = img.shape

arr_label = []  # labels的边界框，存储四个顶点坐标
ori_clas = []  # label的真实类别
for line in lines:
    line = line.split(" ")
    lrbt, c = xywh2lrbt(img_w, img_h, line)
    arr_label.append(lrbt)
    ori_clas.append(c)

# cut_imgs = []  # 被剪切的图像
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 = "./images/cut_imgs/" + str(i) + "_" + str(num) + ".jpg"
    cv2.imwrite(save_cut_ori_img, cut_ori_img)
    save_cut_ori_img = mdFile.new_inline_image(text='', path=save_cut_ori_img)
    num += 1
    ori_cla = ori_clas[i]
    for j in range(len(result[0])):
        res = []
        predict_cla = result[0][j]['class']
        if ori_cla == predict_cla:
            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 = "./images/cut_imgs/" + str(j) + ".jpg"
                cv2.imwrite(save_cut_predict_img, cut_predict_img)
                save_cut_predict_img = mdFile.new_inline_image(text='', path=save_cut_predict_img)
                isTrue = True
                list_item.extend([ori_cla, predict_cla, save_cut_ori_img, save_cut_predict_img, isTrue])
                # list_item.extend(res)
                break
        elif j == len(result[0]) - 1:
            isTrue = False
            list_item.extend([ori_cla, predict_cla, save_cut_ori_img, "", isTrue])
            # res.append([ori_cla, predict_cla, "", "", isTrue])
            list_item.extend(res)
#
# print(list_item)
# print(len(arr_label))
print(len(list_item))
# print(len(result[0]))
mdFile.new_line()
mdFile.new_table(columns=5, rows=len(list_item)//5, text=list_item, text_align='center')
mdFile.create_md_file()

if __name__ == "__main__":
    images_path = "./test_imgs/images/"
    labels_path = "./test_imgs/labels/"
    list_images = os.listdir(images_path)
    list_labels = os.listdir(labels_path)
    for image in list_images:
        img_name = image.split(".")[0]
        mdFile = MdUtils(file_name='yq_0819_' + str(img_name), title='yq_0819_' + str(img_name) + '_iou>=0.5')




#
# images = os.listdir(images_path)
# labels = os.listdir(labels_path)
#
# assert len(images) == len(labels)
#
# for i in range(len(images)):
#     image = images[i]
#     label = labels[i]
#     with open(labels_path + label, mode="r", encoding="utf-8") as f:
#         lines = f.readlines()
#
#     img = images_path + images[i]
#     img = cv2.imread(img)
#     img_h, img_w, _ = img.shape
#
#     arr_label = []  # labels的边界框，存储四个顶点坐标
#     ori_cla = []  # label的真实类别
#     for line in lines:
#         line = line.split(" ")
#         lrbt, c = xywh2lrbt(img_w, img_h, line)
#         arr_label.append(lrbt)
#         ori_cla.append(c)
#
#     cut_path = "./valid_img/cut/"
#
#     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])
#         cut_img = img[lt_y:rb_y, lt_x:rb_x, :]
#         cv2.imwrite(cut_path + str(i) + ".jpg", cut_img)
#
#         box = np.reshape(np.array(arr_label[i]), [-1, 1, 2]).astype(np.int32)
#         img = cv2.polylines(np.array(img), [box], True, (255, 0, 0), 5)
#
#
#
#
# mdFile.create_md_file()
#
#