【目标检测】从txt读取并可视化旋转目标边界框

欢迎光临我的知乎同款文章：https://zhuanlan.zhihu.com/p/396562363

运行环境：

python=3.7
opencv-python==4.5.2.54
opencv-contrib-python==4.5.3.56

代码一：

import os
import os.path as osp
import cv2
import numpy as np
import shutil

BASEDIR = osp.abspath('.')
IMGDIR = osp.join(BASEDIR, 'images')
LABELDIR = osp.join(BASEDIR, 'labelTxt')
OUTDIR = osp.join(BASEDIR, 'watch_GT')
if osp.exists(OUTDIR):
    shutil.rmtree(OUTDIR)
os.makedirs(OUTDIR)


key_list = [name.split('.')[-2] for name in os.listdir(IMGDIR) if name.endswith('.png')]

print(key_list)
for k in key_list:
    imgname = k+'.png' 
    labelname = k+'.txt'
    path_img = osp.join(IMGDIR, imgname)
    path_label = osp.join(LABELDIR, labelname)
    img = cv2.imread(path_img)
    with open(path_label, 'r') as fp:
       lines = fp.readlines()
    # print(lines)
    for line in lines:
        one_box_info = line.strip().split()
        *xy4, cls_name, easy_token = one_box_info
        xy4 = list(map(int, xy4))
        xy4 = [[xy4[0],xy4[1]],
               [xy4[2],xy4[3]],
               [xy4[4],xy4[5]],
               [xy4[6],xy4[7]]]
        xy4 = np.array(xy4)
        # print(xy4)
        cv2.drawContours(img, [xy4], 0, color=(255, 255, 0), thickness=3)
        
    path_out = osp.join(OUTDIR, imgname.split('.')[-2] + '_watchGT.' + imgname.split('.')[-1])
    cv2.imwrite(path_out, img)
    print(imgname, 'done.')

代码一，适用的标注文件格式示例：

（四顶点式标注）

177 345 240 338 250 424 187 431 C 0
356 17 433 4 447 83 370 96 C 0
374 113 448 103 460 186 386 197 C 0
469 202 479 283 398 293 388 212 C 0
400 314 486 304 496 390 410 400 C 0
416 410 494 398 507 480 429 492 C 0

= = = = = = = = = = = 分割线 = = = = = = = = = = = =

代码二：

import cv2
import numpy as np
import os
import os.path as osp


color_dict = {
    0.0: (255, 000, 000),
    1.0: (255, 128, 000),
    2.0: (255, 255, 000),
    3.0: (000, 255, 000),
    4.0: (000, 255, 255),
    5.0: (000, 000, 255),
    6.0: (128, 000, 255),
    7.0: (255, 000, 255),
    8.0: (128, 000, 000),
    9.0: (000, 128, 000),
    10.0: (000, 000, 128)
}


def run_watch_txts(img_dir, txt_dir, out_dir, color):
    IMAGEDIR = osp.abspath(img_dir)
    LABELDIR = osp.abspath(txt_dir)
    OUTDIR = osp.abspath(out_dir)

    png_names = set([name[:-4] for name in os.listdir(IMAGEDIR) if name.endswith('.png')])
    txt_names = set([name[:-4] for name in os.listdir(LABELDIR) if name.endswith('.txt')])
    intersection_names = png_names & txt_names

    for name in png_names:
        image = cv2.imread(osp.join(IMAGEDIR, name+".png"))

        if name not in intersection_names:
            cv2.imwrite(osp.join(OUTDIR, name+".png"), image)
            continue

        IMG_H, IMG_W, IMG_C = image.shape

        with open(osp.join(LABELDIR, name+".txt"), 'r') as fp:
            lines = fp.readlines()

        for line in lines:
            line = line.strip().split(" ")
            # print(line)
            label = float(line[0])
            cx = IMG_W * float(line[1])
            cy = IMG_H * float(line[2])
            w = IMG_W * float(line[3])
            h = IMG_H * float(line[4])
            angle = int(line[5])
            rect = ((cx, cy), (w, h), (angle))
            box = cv2.boxPoints(rect)
            box = np.int0(box)
            cv2.drawContours(image, [box], 0, color_dict[label][::-1], 2)
            # cv2.imshow(" ", image)

        cv2.imwrite(osp.join(OUTDIR, name + ".tif"), image)
        print(name + ", done...")


if __name__ == '__main__':
    BASEDIR = osp.dirname(osp.abspath(__file__))
    
    img_dir = osp.join(BASEDIR, 'imgs')
    txt_dir = osp.join(BASEDIR, 'imgs_txts')
    out_dir = osp.join(BASEDIR, 'imgs_txts_watch')
    if not os.path.exists(out_dir):
        os.makedirs(out_dir)

    run_watch_txts(img_dir, txt_dir, out_dir, color=(250, 200, 250))

代码二，适用的标注文件格式示例：

(浮点型)类别号 中心横坐标cx 中心纵坐标cy 框宽w 框高h 框的角度angle

0.0 0.4375 0.54375 0.040625 0.025 30
0.0 0.56484375 0.5822916666666667 0.0421875 0.03125 45
0.0 0.7671875 0.5572916666666666 0.03125 0.027083333333333334 60

可视化结果示例：

OK！