yolov5常用数据预处理脚本

原创

已于 2023-07-20 19:08:54 修改 · 488 阅读

标签

#YOLO #python #开发语言

收录于

于 2023-07-20 18:56:50 首次发布

该代码主要涉及Python对XML文件的处理，包括修改标签内容、移动文件、删除无标签的XML文件。同时，它实现了一个根据XML文件打Mosaic的效果，特别是针对标注为other的对象进行像素化处理。

文章目录

一、修改xml中标签
二、移动文件以及删除空标签的xml
三、根据xml文件打mosaic

一、修改xml中标签

import os
import shutil
from tqdm import tqdm
import xml.etree.ElementTree as ET
from loguru import logger

@logger.catch()
def datasets_info():
    source = './train/temp'
    instance = {
   
   }
    all_xml = 0
    for root1, dirs, files in os.walk(source):
        for file in files:
            if file.endswith('.xml'):
                all_xml +=1
                path = os.path.join(root1, file)
                tree = ET.parse(path)
                root = tree.getroot()
                objects = root.findall('object')
                for obj in objects:
                    name = obj.find('name').text
                    if name not in instance:
                        instance[name] = 1
                    else:
                        instance[name] += 1

    print('total: %d' % all_xml)
    print('示例框：', instance)

def modify_label():
    source = './train/temp'
    # front_whee working_luggae_transfer_truck
    for root, dirs, files in os.walk(source):
        for file in files:
            fp = os.path.join(root, file)
            if file.endswith('.xml'):
                tree = ET.parse(fp)
                root1 = tree.getroot()
                objects = root1.findall