VOC格式标签与YOLO格式标签可视化

'''
Author: your name
Date: 2021-10-22 17:00:38
LastEditTime: 2021-10-22 17:06:32
LastEditors: Please set LastEditors
Description: In User Settings Edit
FilePath: /mypython/2.vis_xml_yolo.py
'''

'''
   1. VOC格式的xml文件可视化
'''

import xml.etree.ElementTree as ET
import os, cv2

# ------------------------------------------------------
#           可视化一张图片
# ------------------------------------------------------
def read_xml_one(xml_file, img_file):
    
    tree = ET.parse(xml_file)
    root = tree.getroot()
    # imgfile='C:/Users/nansbas/Desktop/01_000002_01244-01086_0939-0989.jpg'
    im = cv2.imread(img_file)
    print(im)
    for object in root.findall('object'):
        object_name = object.find('name').text
        print(object_name)
        Xmin = int(object.find('bndbox').find('xmin').text)
        Ymin = int(object.find('bndbox').find('ymin').text)
        Xmax = int(object.find('bndbox').find('xmax').text)
        Ymax = int(object.find('bndbox').find('ymax').text)
        color = (4, 250, 7)
        cv2.rectangle(im, (Xmin, Ymin), (Xmax, Ymax), color, 2)
        font = cv2.FONT_HERSHEY_SIMPLEX
        cv2.putText(im, object_name, (Xmin, Ymin - 7), font, 0.5, (6, 230, 230), 2)
        # cv2.imshow('01',im)
        print(im)
        cv2.imwrite('./2.jpg', im)


# ---------------------------------------------------
# 可视化文件夹中所有xml文件 #
# ---------------------------------------------------
type45="i2,i4,i5,il100,il60,il80,io,ip,p10,p11,p12,p19,p23,p26,p27,p3,p5,p6,pg,ph4,ph4.5,ph5,pl100,pl120,pl20,pl30,pl40,pl5,pl50,pl60,pl70,pl80,pm20,pm30,pm55,pn,pne,po,pr40,w13,w32,w55,w57,w59,wo"
type45 = type45.split(',')
type45_dic = {}
for item in type45:
    type45_dic[item] = 0
# print(type45_dic)
print(len(type45_dic))
# classes = {'1_23': 0, '3_24_n40': 0, '5_20': 0, '5_16': 0, '5_19_1': 0, '2_1': 0, '2_4': 0, '4_1_1': 0, '3_27': 0, '5_15_2': 0}
def read_xml(ImgPath, AnnoPath, Savepath):
    i = 0
    count = 0
    imagelist = os.listdir(AnnoPath)
    # print(imagelist)
    for image in imagelist:
        image_pre, ext = os.path.splitext(image)
        imgfile = ImgPath + '\\' + image_pre + '.jpg'
        xmlfile = AnnoPath + '\\' + image_pre + '.xml'
        im = cv2.imread(imgfile)
        DomTree = xml.dom.minidom.parse(xmlfile)
        annotation = DomTree.documentElement
        filenamelist = annotation.getElementsByTagName('filename')
        filename = filenamelist[0].childNodes[0].data
        objectlist = annotation.getElementsByTagName('object')

        for objects in objectlist:
            count += 1

            namelist = objects.getElementsByTagName('name')
            objectname = namelist[0].childNodes[0].data
            # if objectname not in classes:
            #     classes[objectname] = 0
            if objectname not in type45:
                # type45[objectname] = 0
                continue
            type45_dic[objectname] += 1

            bndbox = objects.getElementsByTagName('bndbox')

            for box in bndbox:
                i += 1
                try:
                    x1_list = box.getElementsByTagName('xmin')
                    x1 = round(float(x1_list[0].childNodes[0].data))
                    y1_list = box.getElementsByTagName('ymin')
                    y1 = round(float(y1_list[0].childNodes[0].data))
                    x2_list = box.getElementsByTagName('xmax')
                    x2 = round(float(x2_list[0].childNodes[0].data))
                    y2_list = box.getElementsByTagName('ymax')
                    y2 = round(float(y2_list[0].childNodes[0].data))
                    #
                    minX = x1
                    minY = y1
                    maxX = x2
                    maxY = y2

                    if (i % 9 == 0):
                        color = (128, 0, 0)
                    elif (i % 9 == 1):
                        color = (153, 51, 0)
                    elif (i % 9 == 2):
                        color = (255, 204, 0)
                    elif (i % 9 == 3):
                        color = (0, 51, 0)
                    elif (i % 9 == 4):
                        color = (51, 204, 204)
                    elif (i % 9 == 5):
                        color = (128, 0, 128)
                    elif (i % 9 == 6):
                        color = (0, 255, 255)
                    elif (i % 9 == 7):
                        color = (60, 179, 113)
                    elif (i % 9 == 8):
                        color = (255, 127, 80)
                    elif (i % 9 == 9):
                        color = (0, 255, 0)
                    cv2.rectangle(im, (x1, y1), (x2, y2), color, 2)
                    # path = Savepath + '/' + image_pre + '.png'
                    # print(path)
                    font = cv2.FONT_HERSHEY_SIMPLEX
                    cv2.putText(im, objectname, (minX, minY - 7), font, 0.5, color, 2)
                    cv2.imwrite(path_to_save, im)
                except Exception as e:
                    print(e)
    print(type45_dic)

if __name__ == '__main__':
    count = 0
    # read_xml_one('I:\Jupyter_file\data\\xmlLabel\\train\\23.xml', 'I:\Jupyter_file\data\\train\\23.jpg')
    read_xml('I:\data\VOCdevkit\VOC2007\JPEGImages', 'I:\data\VOCdevkit\VOC2007\Annotations',
             'I:\data\\vis_img')



'''
    2.YOLO格式的txt标签文件可视化
'''

import cv2
import os

def draw_box_in_single_image(image_path, txt_path):
    # 读取图像
    image = cv2.imread(image_path)

    # 读取txt文件信息
    def read_list(txt_path):
        pos = []
        with open(txt_path, 'r') as file_to_read:
            while True:
                lines = file_to_read.readline()  # 整行读取数据
                if not lines:
                    break
                # 将整行数据分割处理，如果分割符是空格，括号里就不用传入参数，如果是逗号， 则传入‘，'字符。
                p_tmp = [float(i) for i in lines.split(' ')]
                pos.append(p_tmp)  # 添加新读取的数据
                # Efield.append(E_tmp)
                pass
        return pos


    # txt转换为box
    def convert(size, box):
        xmin = (box[1]-box[3]/2.)*size[1]
        xmax = (box[1]+box[3]/2.)*size[1]
        ymin = (box[2]-box[4]/2.)*size[0]
        ymax = (box[2]+box[4]/2.)*size[0]
        box = (int(xmin), int(ymin), int(xmax), int(ymax))
        return box

    pos = read_list(txt_path)
    tl = int((image.shape[0]+image.shape[1])/2) +1
    lf = max(tl-1,1)
    for i in range(len(pos)):
        print(i)
        label = str(int(pos[i][0]))
        box = convert(image.shape, pos[i])
        image = cv2.rectangle(image,(box[0], box[1]),(box[2],box[3]),(0,0,255),2)
        cv2.putText(image,label,(box[0],box[1]-2), 0, tl, [0,0,255], thickness=1, lineType=cv2.LINE_AA)
        pass

    cv2.imwrite('I:\data\\vis_img\\{}.png'.format(image_path.split('\\')[-1][:-4]), image)
    # cv2.imshow("images", image)
    # cv2.waitKey(0)
    # cv2.destroyAllWindows()


img_folder = "I:\data\VOCdevkit\VOC2007\JPEGImages\\"
img_list = os.listdir(img_folder)
img_list.sort()

label_folder = "I:\data\VOCdevkit\VOC2007\labels\\"
label_list = os.listdir(label_folder)
label_list.sort()

for i in range(len(img_list)):
    image_path = img_folder + "\\" + img_list[i]
    txt_path = label_folder + "\\" + label_list[i]
    draw_box_in_single_image(image_path, txt_path)