[我参加NVIDIA Sky Hackathon]（数据处理脚本）

由于比赛需要自己收集数据并打标签，标签规定使用kitti格式。由于labelme标注工具默认生成的时voc格式， 所以需要把voc格式转化成kitti格式，所以写了一个voc转kitti的脚本，除此之外，为了验证打的标签是否准确，还写了一个验证脚本，都是用python写的。

1.voc转kitti脚本

# 作 者:Leo
# 时 间:2022/11/15 19:12
import xml.etree.ElementTree as ET
import os
from tqdm import tqdm
import cv2
'''
base_xml_dir为xml文件夹路径
kitti_saved_dir为要保存的kitti文件夹的路径，转化前，先把输出文件夹创建好
'''
base_xml_dir = "E:\\PyProject\\tools\\ba\\a\\xml\\"
xml_list = os.listdir(base_xml_dir)
kitti_saved_dir = "E:\\PyProject\\tools\\ba\\a\\kitti\\"

clslist = []
def convert_annotation(file_name):
    #print('********')
    #print(base_xml_dir,file_name)
    in_file = open(base_xml_dir + file_name, 'rb')
    tree = ET.parse(in_file)
    root = tree.getroot()

    with open(kitti_saved_dir + file_name[:-4] + '.txt', 'w') as f:
        for obj in root.iter('object'):
            cls = obj.find('name').text
            # if cls not in object_class_list:
            #     object_class_list.append(cls)
            xmlbox = obj.find('bndbox')
            """
                第5～8这4个数：物体的2维边界框
                xmin，ymin，xmax，ymax
            """
            xmin, ymin, xmax, ymax = xmlbox.find('xmin').text, xmlbox.find('ymin').text, \
                                     xmlbox.find('xmax').text, xmlbox.find('ymax').text
            a = int(xmin)
            b = int(ymin)
            c = int(xmax)
            d = int(ymax)
            #print(xmin, ymin, xmax, ymax)
            #print(type(float(xmin)))
            #cv2.rectangle(image, (xmlbox.find('xmin').text, xmlbox.find('ymin').text), (xmlbox.find('xmax').text, xmlbox.find('ymax').text), (0, 0, 255), 2)
            #cv2.rectangle(image, (0, 55), (224, 263), (0, 0, 255), 2)
            #print(cls)
            if cls == 'Peel_the_fruit':
                #print('******1*******')
                #print(cls)
                cls = 'banane'
                #cv2.rectangle(image, (a, b), (c, d), (0, 0, 255), 2)
                #print('香蕉皮')
                #print(cls)
                #print('******2*******')
            elif cls == 'pop_top_can':
                #print(cls)
                cls = 'bottle'
                #cv2.rectangle(image, (a, b), (c, d), (0, 0, 255), 2)
                #print('瓶子1')
            elif cls == 'drink_bottle':
                cls = 'bottle'
                #cv2.rectangle(image, (a, b), (c, d), (0, 0, 255), 2)
                #print('瓶子2')
            elif cls == 'carton_carton':
                #print(cls)
                cls = 'CARDBOARD'
                #cv2.rectangle(image, (a, b), (c, d), (0, 0, 255), 2)
                #print('纸箱')
            elif cls == 'Drink_box':
                cls = 'CARDBOARD'
                #cv2.rectangle(image, (a, b), (c, d), (0, 0, 255), 2)
            else:
                pass
            if cls == 'banane' or cls == 'bottle' or cls == 'CARDBOARD':
                f.write(cls + " " + '0.00' + " " + '0' + " " + '0.0' + " " + str(xmin) + " "
                        + str(ymin) + " " + str(xmax) + " " + str(ymax) + " " +
                        '0.0' + " " + '0.0' + " " + '0.0' + " " + '0.0' + " " + '0.0' + " " + '0.0' + " " + '0.0' + '\n')

# object_class_list = []  # ['trafficsignal', 'car', 'person', 'trafficlight', 'bicycle', 'motorbike', 'bus']

for i in tqdm(range(len(xml_list))):
    #image_name = xml_list[i].replace('xml', 'jpg')
    #image = cv2.imread(image_name)
    #cv2.rectangle(image, (0, 55), (224, 263), (0, 0, 255), 2)
    convert_annotation(xml_list[i])
    #cv2.imwrite(image_name.replace('.jpg', '-box.jpg'), image)
    #print('原始',xml_list[i])
    #print('最后', image_name, xml_list[i])
    # print(i)
#print(clslist)
# print(object_class_list)

                                                效果如下图1

 

 

2.验证打框脚本

# 作 者:Leo
# 时 间:2022/11/21 21:14
import os
from tqdm import tqdm
import cv2
'''
base_img_dir为未打框图片的文件夹路径
base_label_dir为kitti标签格式的文件夹路径
op_img_dir为打框后的图片输出的文件夹路径，需要在输出钱创建好
'''
base_img_dir = "E:\\PyProject\\tools\\t1\\image_comb\\"
base_label_dir = "E:\\PyProject\\tools\\t1\\kitti_img\\"
op_img_dir = "E:\\PyProject\\tools\\t1\\img-box\\"
img_list = os.listdir(base_img_dir)
label_list = os.listdir(base_label_dir)
def draw_imgbox(image, label):
	for line in open(label, "r", encoding='UTF-8'):
		line_list = []
		line_list = line.split(' ')
		line_list = line_list[0:1] + line_list[4:8]
		print(line_list)
		cv2.rectangle(img_file, (int(line_list[1]), int(line_list[2])), (int(line_list[3]), int(line_list[4])), (0, 0, 255), 2)

for i in tqdm(range(len(img_list))):
	print('第' + str(i) + '张')
	img_file = cv2.imread(base_img_dir + img_list[i])
	label_file = label_list[i]
	draw_imgbox(img_file, base_label_dir + label_file)
	cv2.imwrite(op_img_dir + img_list[i].replace('.jpg', '-box.jpg'), img_file)

3.两个脚本效果图如下