echodict/ppv5

1 2import json3import glob4from pathlib import Path5import cv26import numpy as np7 8def draw_box(pth_webp, pth_boxs):9 10    with open(pth_boxs, encoding='utf-8') as fp:11        boxs = json.load(fp)12        word_boxes = boxs['text_word_boxes']13        word_boxes = word_boxes[::-1] # 逆序14        text_word = boxs['text_word']15        text_word = text_word[::-1]16        rec_texts = boxs['rec_texts']17        rec_texts = rec_texts[::-1]18        fp.close()19 20    imgData = np.fromfile(pth_webp, dtype=np.uint8)21    img = cv2.imdecode(imgData, cv2.IMREAD_UNCHANGED)22 23    if len(img.shape) != 3:  # 转彩图24        img_color = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)25        img = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)  # DBNet 原版只能处理彩图，这里转一下26    else:27        img_color = img.copy()28 29 30    for arr_word, arr_box in zip( text_word, word_boxes):31        for word, box in zip(arr_word, arr_box):32            lu = [box[0], box[1]]33            rd = [box[2], box[3]]34            ld = [box[0], box[3]]35            ru = [box[2], box[1]]36            points = np.array([lu, ru, rd, ld])37 38            cv2.polylines(img_color, [points], isClosed=True, color=(  # 多边形，框得比较全39                100, 0, 255), thickness=2)  # 只画线，不填充40            41            cv2.imshow("box", img_color)42            cv2.waitKey(0)43 44    path_jpg = pth_webp.replace(".png", ".jpg")45    cv2.imwrite(path_jpg, img_color)46 47def do_drawbox():48 49    pth_box = 'output/SWX0005_00000_00001_res.json'50    pth_webp = 'output/SWX0005_00000_00001_input_img.png'51    draw_box(pth_webp, pth_box)52 53if __name__ == "__main__":54    do_drawbox()55