常用脚本-持续更新(文件重命名、视频抽帧、拆帧、删除冗余文件、yolo2xml、转换图片格式、修改xml)

所有代码位置Learning-Notebook-Codes/Python/常用脚本

1. 文件重命名

  • 脚本路径:codes/files_rename.py
  • 脚本说明:可以自动重命名某个文件夹下指定类型的文件。
    • 修改前文件名称: img1.jpg
    • 修改后文件名称: Le0v1n-20231123-X-0001.jpg
python 复制代码
import os
import tqdm
import datetime


"""============================ 需要修改的地方 ==================================="""
SRC_PATH = 'Python/常用脚本/EXAMPLE_FOLDER'  # 文件夹路径
file_type = ('.png', '.jpg', '.jpeg', '.gif')  # 想要重命名的文件类型

# -------------------重命名相关------------------
retain_previous_name = False  # 是否保留之前的名称
new_name = "Le0v1n"  # retain_previous_name为False时生效
use_date_stamp = True  # 是否使用时间戳 -> e.g. 20231123
comment = "X"  # 备注
use_serial_numbering = True  # 是否使用顺序的编号 -> 1, 2, 3, 4, 5, 6, ...
start_number = 1  # 从编号几开始 -> e.g. 1: 从 0001 开始编号
numbering_placeholder = 4  # 编号保留的占位 -> e.g. 0001, 0002, 0003, ...
hyphen = '-'  # 连字符 -> e.g. filename-0001.jpg
"""==============================================================================="""

# 获取目录中的所有图片文件
files_list = [file for file in os.listdir(SRC_PATH) if file.lower().endswith(file_type)]

"------------计数------------"
TOTAL_FILES_NUM = len(files_list)  # 需要重命名的文件数量
RENAME_NUM = 0  # 重命名成功数量
"---------------------------"

# 获取当前时间并格式化时间戳
timestamp = datetime.datetime.now().strftime("%Y%m%d")

# 遍历文件
process_bar = tqdm.tqdm(total=TOTAL_FILES_NUM, desc="为指定格式的文件重命名", unit='file')  # 创建进度条
for idx, file_name in enumerate(files_list):
    file_pre, file_ext = os.path.splitext(file_name)  # 获得文件名和后缀
    process_bar.set_description(f"rename for \033[1;31m{file_name}\033[0m")

    # 构建新的文件名
    if retain_previous_name:  # 保留原有的名称
        NEW_FILE_NAME = f"{file_pre}"
    elif new_name:  # 不保留原有的名称且新名称存在
        NEW_FILE_NAME = new_name
    else:  # 不保留原有的名称也没有新名称 -> 报错
        raise KeyError(f"不保留原有的名称也没有新名称!")
    
    if use_date_stamp:  # 使用时间戳
        NEW_FILE_NAME += f"{hyphen}{timestamp}"
    
    if comment:  # 添加备注
        NEW_FILE_NAME += f"{hyphen}{comment}" 
    
    if use_serial_numbering:  # 使用编号
        NEW_FILE_NAME += f"{hyphen}{idx + start_number:0{numbering_placeholder}d}"

    # 加上扩展名
    NEW_FILE_NAME += file_ext
    
    # 开始重命名文件         
    _src = os.path.join(SRC_PATH, file_name)  # 旧文件路径
    _dst = os.path.join(SRC_PATH, NEW_FILE_NAME)  # 新文件路径
    
    os.rename(_src, _dst)  # 重命名文件
    RENAME_NUM += 1
    process_bar.update(1)
process_bar.close()
    
print(f"👌 文件重命名完成: {RENAME_NUM}/{TOTAL_FILES_NUM}")

2. 视频抽帧

  • 脚本路径:codes/extract_frames.py
  • 脚本说明:根据帧间隔对某个文件夹下指定类型的视频文件进行抽帧,得到系列图片。
    • 视频文件所在文件夹名称: EXAMPLE_FOLDER
    • 抽帧得到的文件夹名称: EXAMPLE_FOLDER/extract_frames_results/test_vid_0001.jpg
python 复制代码
import cv2
import os
import tqdm
from utils import create_folder


"""============================ 需要修改的地方 ==================================="""
SRC_PATH = "Python/常用脚本/EXAMPLE_FOLDER"  # 原始视频路径
frame_interval = 10  # 视频采样间隔,越小采样率越高 -> 60 | 30 | 15 | 10
video_type = ['.mp4', '.avi']  # 视频格式(.mp4 | .avi)

DST_PATH = "extract_frames_results"  # 保存图片文件夹名称
save_img_format = '.jpg'  # 保存的图片格式(.jpg | .png)
"""==============================================================================="""

# 构建路径
results_imgs_path = os.path.join(SRC_PATH, DST_PATH)  # 保存图片路径

# 得到存放所有视频的list
video_list = [x for x in os.listdir(SRC_PATH) if os.path.splitext(x)[-1] in video_type]

"------------计数------------"
TOTAL_VID_NUM = len(video_list)
SUCCEED_NUM = 0  # 完成视频的个数
TOTAL_IMG_NUM = 0  # 统计得到的所有图片数量
"---------------------------"

print(f"\033[1;31m[SRC]视频路径为: {SRC_PATH}\033[0m"
      f"\n\t\033[1;32m视频个数: {TOTAL_VID_NUM}\033[0m"
      f"\n\033[1;31m[DST]图片保存路径为: {DST_PATH}\033[0m"
      f"\n\t\033[1;32m保存的图片格式为: {save_img_format}\033[0m"
      f"\n\n请输入 \033[1;31m'yes'\033[0m 继续,输入其他停止")
_INPUT = input()
if _INPUT != "yes":
    exit()
    
create_folder(results_imgs_path, verbose=True)  # 创建文件夹

# 创建一个tqdm进度条对象
progress_bar = tqdm.tqdm(total=len(video_list), desc="视频拆帧...", unit="vid")
statistics_dict = dict()  # 创建一个字典,用于统计
for vid_name in video_list:  # 遍历所有的视频
    save_number = 1  # 记录当前视频保存的frame个数
    vid_pre, vid_ext = os.path.splitext(vid_name)  # 获取文件名和后缀
    
    vid_path = os.path.join(SRC_PATH, vid_name)  # 视频完整路径
    
    # 创建VideoCapture对象
    vc = cv2.VideoCapture(vid_path)

    # 检查视频是否成功打开
    if not vc.isOpened():
        continue
    
    # 逐帧读取视频并保存为图片
    frame_count = 0
    while True:
        # 读取一帧
        rval, frame = vc.read()

        # 检查是否成功读取帧
        if not rval:  # 读取帧失败
            break

        # 每隔 frame_interval 帧保存一次图片
        if frame_count % frame_interval == 0:
            # 生成图片文件名
            frame_name = f"{vid_pre}_{save_number:04d}{save_img_format}"
            frame_path = os.path.join(results_imgs_path, frame_name)  # Python\常用脚本\EXAMPLE_FOLDER\extract_frames_results\test_vid_0016.jpg

            progress_bar.set_description(f"\033[1;31m{vid_name}\033[0m -> "
                                            f"\033[1;36m{save_number * frame_interval:04d}\033[0m"
                                            f" ({save_number})")  # 更新tqdm的描述
            # 保存帧为图片
            cv2.imwrite(frame_path, frame)
            save_number += 1

        # 帧数加1
        frame_count += 1

    # 释放VideoCapture对象
    vc.release()
    TOTAL_IMG_NUM += save_number  # 更新图片数量
    SUCCEED_NUM += 1
    statistics_dict[vid_pre] = save_number  # 更新字典,记录当前视频得到的frame个数
    progress_bar.update()  
progress_bar.close()

print("------------------------------------------------------------------")
_cont = 0
for k, v in statistics_dict.items():
    print(f"\033[1;34m"
          f"👌 1. [{k}] 得到 frame 个数 -> {v}"
          f"\033[0m")
    _cont += 1
print()
print(f"\033[1;31m"
      f"👌👌👌 视频拆帧 ({TOTAL_VID_NUM}个)完成,总共得到[{TOTAL_IMG_NUM}]张{save_img_format}图片!"
      f"\033[0m")
print("------------------------------------------------------------------")

3. 根据文件A删除冗余的文件B

  • 脚本路径:codes/delete-Redundant_fileB.py
  • 脚本说明:根据文件 A 删除冗余文件 B。
  • 用途:根据 annotations 删除冗余的 images。
  • 要求:文件 A 和 文件 B 应该有相同的名字(后缀不同)。
python 复制代码
"""
+ 脚本说明:根据文件 A 删除冗余文件 B。
+ 用途:根据 annotations 删除冗余的 images。
+ 要求:文件 A 和 文件 B 应该有相同的名字(后缀不同)。
"""
import os
import tqdm


"""============================ 需要修改的地方 ==================================="""
path_A = 'Python/常用脚本/EXAMPLE_FOLDER/images'  # 不删除
file_type_A = ('.jpg', '.png')

path_B = 'Python/常用脚本/EXAMPLE_FOLDER/annotations'  # 会删除的
file_type_B = ('.json', '.xml')
"""==============================================================================="""

# 获取两种文件列表
files_A_list = [file for file in os.listdir(path_A) if file.endswith(file_type_A)]
files_B_list = [file for file in os.listdir(path_B) if file.endswith(file_type_B)]

"------------计数------------"
NUM_B = len(files_A_list)
NUM_B = len(files_B_list)
SUCCEED_NUM = 0
SKIP_NUM = 0
"---------------------------"

print(f"文件[A]所在文件夹路径为: {path_A}"
      f"\n\t文件[A]数量为: {NUM_B}"
      f"\n\t文件[A]的后缀为: {file_type_A}"
      f"\n文件[B]所在文件夹路径为: {path_B}"
      f"\n\t文件[B]数量为: {NUM_B}"
      f"\n\t文件[B]的后缀为: {file_type_B}"
      f"\n\n请输入 \033[1;31m'yes'\033[0m 继续,输入其他停止")
_INPUT = input()
if _INPUT != "yes":
    exit()

# 遍历文件B
process_bar = tqdm.tqdm(total=NUM_B, desc="根据文件A删除冗余的文件B", unit='unit')
for name_B in files_B_list:
    pre_B, ext_A = os.path.splitext(name_B)  # 分离文件名和后缀
    process_bar.set_description(f"Process with \033[1;31m{name_B}\033[0m")
    
    # 判断对应的同名 A 文件是否存在,如果存在则跳过
    dst_path = os.path.join(path_A, pre_B)  # 没有后缀
    _exist_flag = 0
    for ext_A in file_type_A:  # 遍历所有格式,看是否有至少一个同名文件存在
        if os.path.exists(dst_path + ext_A):
            _exist_flag += 1
    if _exist_flag > 0:  # 如果存在至少一个同名文件, 则跳过
        SKIP_NUM += 1
        process_bar.update()
    else:  # 没有同名文件, 则删除文件B
        del_path = os.path.join(path_B, name_B)
        os.remove(del_path)
        SUCCEED_NUM += 1
        process_bar.update()
process_bar.close()

print(f"👌 冗余的B文件删除已完成!"
      f"\n\t删除文件数量/文件B数量 = {SUCCEED_NUM}/{NUM_B}"
      f"\n\t跳过文件数量/文件B数量 = {SKIP_NUM}/{NUM_B}")

if SUCCEED_NUM + SKIP_NUM == NUM_B:
    print("👌 No Problems")
else:
    print(f"🤡 有问题,请仔细核对!"
          f"\n\tSUCCEED_NUM: {SUCCEED_NUM}\tSKIP_NUM: {SKIP_NUM}"
          f"\n\tSUCCEED_NUM + SKIP_NUM + ERROR_NUM = {SUCCEED_NUM + SKIP_NUM}"
          f"\n\tTOTAL_NUM: {NUM_B}")

4. yolo2xml

  • 脚本路径:codes/D-yolo2xml.py
  • 脚本说明:将yolo格式txt标注文件转换为voc格式xml标注文件
  • 用途:将 YOLO 格式的标签文件还原为 xml 格式
  • 要求:图片和yolo标签应该有相同的名字(后缀不同)
python 复制代码
"""
+ 脚本说明:将yolo格式txt标注文件转换为voc格式xml标注文件
+ 用途:将 YOLO 格式的标签文件还原为 xml 格式
+ 要求:图片和yolo标签应该有相同的名字(后缀不同)
"""
from xml.dom.minidom import Document
import os
import cv2
import tqdm


"""============================ 需要修改的地方 ==================================="""
IMAGE_PATH = "EXAMPLE_FOLDER/images"  # 原图文件夹路径
TXT_PATH = "EXAMPLE_FOLDER/labels-yolo"  # 原txt标签文件夹路径
XML_PATH = "EXAMPLE_FOLDER/labels-xml"  # 保存xml文件夹路径
image_type = '.jpg'
create_empty_xml_for_neg = True  # 是否为负样本生成对应的空的xml文件


classes_dict = {
    '0': "cat",
    '1': 'dog'
}
"""==============================================================================="""

os.makedirs(XML_PATH) if not os.path.exists(XML_PATH) else None

txt_file_list = [file for file in os.listdir(TXT_PATH) if file.endswith(".txt") and file != 'classes.txt']

"------------计数------------"
TOTAL_NUM = len(txt_file_list)
SUCCEED_NUM = 0  # 成功创建xml数量
SKIP_NUM = 0  # 跳过创建xml文件数量
OBJECT_NUM = 0  # object数量
"---------------------------"

process_bar = tqdm.tqdm(total=TOTAL_NUM, desc="yolo2xml", unit='.txt')
for i, txt_name in enumerate(txt_file_list):
    process_bar.set_description(f"Process in \033[1;31m{txt_name}\033[0m")
    txt_pre, txt_ext = os.path.splitext(txt_name)  # 分离前缀和后缀
    
    xmlBuilder = Document()  # 创建一个 XML 文档构建器
    annotation = xmlBuilder.createElement("annotation")  # 创建annotation标签
    xmlBuilder.appendChild(annotation)
    
    # 打开 txt 文件
    txtFile = open(os.path.join(TXT_PATH, txt_name))
    txtList = txtFile.readlines()  # 以一行的形式读取txt所有内容
    
    if not txtList and not create_empty_xml_for_neg:  # 如果 txt 文件内容为空且不允许为负样本创建xml文件
        SKIP_NUM += 1
        process_bar.update()
        continue
        
    # 读取图片
    img = cv2.imread(os.path.join(IMAGE_PATH, txt_pre) + image_type)
    H, W, C = img.shape
    
    # folder标签
    folder = xmlBuilder.createElement("folder")  
    foldercontent = xmlBuilder.createTextNode('images')
    folder.appendChild(foldercontent)
    annotation.appendChild(folder)  # folder标签结束

    # filename标签
    filename = xmlBuilder.createElement("filename")  
    filenamecontent = xmlBuilder.createTextNode(txt_pre + image_type)
    filename.appendChild(filenamecontent)
    annotation.appendChild(filename)  # filename标签结束

    # size标签
    size = xmlBuilder.createElement("size")  
    width = xmlBuilder.createElement("width")  # size子标签width
    widthcontent = xmlBuilder.createTextNode(str(W))
    width.appendChild(widthcontent)
    size.appendChild(width)  # size子标签width结束

    height = xmlBuilder.createElement("height")  # size子标签height
    heightcontent = xmlBuilder.createTextNode(str(H))
    height.appendChild(heightcontent)
    size.appendChild(height)  # size子标签height结束

    depth = xmlBuilder.createElement("depth")  # size子标签depth
    depthcontent = xmlBuilder.createTextNode(str(C))
    depth.appendChild(depthcontent)
    size.appendChild(depth)  # size子标签depth结束
    annotation.appendChild(size)  # size标签结束
    
    # 读取 txt 内容,生成 xml 文件内容
    for line in txtList:  # 正样本(txt内容不为空)
        # .strip()去除行首和行尾的空白字符(如空格和换行符)
        oneline = line.strip().split(" ")  # oneline是一个list, e.g. ['0', '0.31188484251968507', '0.6746135899679205', '0.028297244094488208', '0.04738990959463407']

        # 开始 object 标签
        object = xmlBuilder.createElement("object")  # object 标签
        
        # 1. name标签
        picname = xmlBuilder.createElement("name")  
        namecontent = xmlBuilder.createTextNode(classes_dict[oneline[0]])  # 确定是哪个类别
        picname.appendChild(namecontent)
        object.appendChild(picname)  # name标签结束

        # 2. pose标签
        pose = xmlBuilder.createElement("pose")  
        posecontent = xmlBuilder.createTextNode("Unspecified")
        pose.appendChild(posecontent)
        object.appendChild(pose)  # pose标签结束

        # 3. truncated标签
        truncated = xmlBuilder.createElement("truncated")  
        truncatedContent = xmlBuilder.createTextNode("0")
        truncated.appendChild(truncatedContent)
        object.appendChild(truncated)  # truncated标签结束
        
        # 4. difficult标签
        difficult = xmlBuilder.createElement("difficult")  
        difficultcontent = xmlBuilder.createTextNode("0")
        difficult.appendChild(difficultcontent)
        object.appendChild(difficult)  # difficult标签结束

        # 5. bndbox标签
        bndbox = xmlBuilder.createElement("bndbox")  
        ## 5.1 xmin标签
        xmin = xmlBuilder.createElement("xmin")  
        mathData = int(((float(oneline[1])) * W + 1) - (float(oneline[3])) * 0.5 * W)
        xminContent = xmlBuilder.createTextNode(str(mathData))
        xmin.appendChild(xminContent)
        bndbox.appendChild(xmin)  # xmin标签结束

        ## 5.2 ymin标签
        ymin = xmlBuilder.createElement("ymin")  # ymin标签
        mathData = int(((float(oneline[2])) * H + 1) - (float(oneline[4])) * 0.5 * H)
        yminContent = xmlBuilder.createTextNode(str(mathData))
        ymin.appendChild(yminContent)
        bndbox.appendChild(ymin)  # ymin标签结束
        
        ## 5.3 xmax标签
        xmax = xmlBuilder.createElement("xmax")  # xmax标签
        mathData = int(((float(oneline[1])) * W + 1) + (float(oneline[3])) * 0.5 * W)
        xmaxContent = xmlBuilder.createTextNode(str(mathData))
        xmax.appendChild(xmaxContent)
        bndbox.appendChild(xmax)  # xmax标签结束

        ## 5.4 ymax标签
        ymax = xmlBuilder.createElement("ymax")  # ymax标签
        mathData = int(
            ((float(oneline[2])) * H + 1) + (float(oneline[4])) * 0.5 * H)
        ymaxContent = xmlBuilder.createTextNode(str(mathData))
        ymax.appendChild(ymaxContent)
        bndbox.appendChild(ymax)  # ymax标签结束

        object.appendChild(bndbox)  # bndbox标签结束
        annotation.appendChild(object)  # object标签结束
        
        OBJECT_NUM += 1

    # 创建 xml 文件
    f = open(os.path.join(XML_PATH, txt_pre) + '.xml', 'w')

    # 为 创建好的 xml 文件写入内容
    xmlBuilder.writexml(f, indent='\t', newl='\n',
                        addindent='\t', encoding='utf-8')
    f.close()  # 关闭xml文件
    
    SUCCEED_NUM += 1
    process_bar.update()
process_bar.close()

print(f"👌yolo2xml已完成, 详情如下:"
      f"\n\t成功转换文件数量/总文件数量 = \033[1;32m{SUCCEED_NUM}\033[0m/{TOTAL_NUM}"
      f"\n\t跳过转换文件数量/总文件数量 = \033[1;31m{SKIP_NUM}\033[0m/{TOTAL_NUM}"
      f"\n\t所有样本的 object 数量/总文件数量 = \033[1;32m{OBJECT_NUM}\033[0m/{TOTAL_NUM}"
      f"\n\t平均每个xml文件中object的数量为: {int(OBJECT_NUM / SUCCEED_NUM)}")

if SUCCEED_NUM + SKIP_NUM == TOTAL_NUM:
    print(f"\n👌 \033[1;32mNo Problem\033[0m")
else:
    print(f"\n🤡 \033[1;31m貌似有点问题, 请仔细核查!\033[0m")

5. 转换图片格式

  • 脚本路径:codes/E-转换图片格式
  • 脚本说明:对指定文件夹下所有的图片进行格式转换
  • 用途:统一数据集图片的格式
  • 要求:无
  • 注意:
    1. 不需要转换的则跳过
    2. 不是图片的文件会扔到指定位置 RECYCLE_BIN_PATH
python 复制代码
"""
+ 脚本说明:对指定文件夹下所有的图片进行格式转换
+ 用途:统一数据集图片的格式
+ 要求:无
+ 注意:
  1. 不需要转换的则跳过
  2. 不是图片的文件会扔到指定位置 RECYCLE_BIN_PATH
"""
import os
import tqdm
from PIL import Image
import shutil


"""============================ 需要修改的地方 ==================================="""
# 定义文件夹路径
IMG_PATH                 = "EXAMPLE_FOLDER/images"  # 输入图片所在文件夹路径
wanna_convert_image_type = '.jpg'  # 想要转换的图片格式
other_image_type         = ['.png', '.jpeg']  # 什么格式的图片将会被转换
"""==============================================================================="""

# 确定回收站位置
RECYCLE_BIN_PATH = os.path.join(os.path.dirname(IMG_PATH), "recycle_bin")

# 获取文件夹内所有文件
all_files = os.listdir(IMG_PATH)

"------------计数------------"
TOTAL_NUM           = len(all_files)
SUCCEED_CONVERT_NUM = 0
SKIP_CONVERT_NUM    = 0
OTHER_FILE_NUM      = 0
"---------------------------"

# 遍历所有的图片
process_bar = tqdm.tqdm(total=TOTAL_NUM, desc=f"将所有图片转换为{wanna_convert_image_type}格式", unit='file')
for file_name in all_files:
    # 分离文件名和后缀
    file_pre, file_ext = os.path.splitext(file_name)
    process_bar.set_description(f"Process in \033[1;31m{file_name}\033[0m")
    
    # 构建文件完整路径
    file_path = os.path.join(IMG_PATH, file_name)
    
    # 检查文件是否为.jpg格式
    if file_ext == wanna_convert_image_type:  # 如果是 jpg 则跳过
        SKIP_CONVERT_NUM += 1
        process_bar.update()
        continue
    elif file_ext in other_image_type:  # 如果是其他图片格式
        with Image.open(file_path) as img:
            # 构建输出文件路径
            dst_save_path = os.path.join(IMG_PATH, file_pre) + wanna_convert_image_type
            img.save(dst_save_path)  # 保存为.jpg格式
            
            # 将原有的图片移动到其他文件夹下
            dst_move_path = os.path.join(RECYCLE_BIN_PATH, file_name)
            shutil.move(src=file_path, dst=dst_move_path)

            SUCCEED_CONVERT_NUM += 1
            process_bar.update()
    else:  # 既不是 jpg 也不是 png、jpeg,则移动到其他文件夹下
        if not os.path.exists(RECYCLE_BIN_PATH):
            os.mkdir(RECYCLE_BIN_PATH)
            
        dst_move_path = os.path.join(RECYCLE_BIN_PATH, file_name)
        shutil.move(src=file_path, dst=dst_move_path)
        OTHER_FILE_NUM += 1
        process_bar.update()
process_bar.close()

print(f"👌 所有图片已转换为jpg, 详情如下:"
      f"\n\t成功转换数量/总文件数量 = \033[1;32m{SUCCEED_CONVERT_NUM}\033[0m/{TOTAL_NUM}"
      f"\n\t跳过文件数量/总文件数量 = \033[1;34m{SKIP_CONVERT_NUM}\033[0m/{TOTAL_NUM}"
      f"\n\t其他格式文件数量/总文件数量 = \033[1;31m{OTHER_FILE_NUM}\033[0m/{TOTAL_NUM}")

if SUCCEED_CONVERT_NUM + SKIP_CONVERT_NUM + OTHER_FILE_NUM == TOTAL_NUM:
    print("👌 No Problems")
else:
    print(f"🤡 貌似有点问题, 请仔细核查!"
          f"\n\tSUCCEED_NUM: {SUCCEED_CONVERT_NUM}"
          f"\n\tSKIP_NUM: {SKIP_CONVERT_NUM}"
          f"\n\tOTHER_FILE_NUM = {OTHER_FILE_NUM}"
          f"\nSUCCEED_NUM + SKIP_NUM + OTHER_FILE_NUM = {SUCCEED_CONVERT_NUM + SKIP_CONVERT_NUM + OTHER_FILE_NUM}"
          f"\nTOTAL_NUM: {TOTAL_NUM}")

6. 根据图片修改xml文件中的size尺寸信息

  • 脚本路径:codes/F-根据图片修改xml文件中的size尺寸信息.py
  • 脚本说明:根据图片修改xml文件中的size尺寸信息
  • 用途:修正数据集标签的信息
  • 要求:无
  • 注意:
    1. 不是in-place操作
    2. 不需要转换的也会复制到新的文件夹下
    3. 如果遇到xml没有对应图片的,则会记录该错误,并生成 ERROR_LOG.txt 文件
python 复制代码
"""
+ 脚本说明:根据图片修改xml文件中的size尺寸信息
+ 用途:修正数据集标签的<size>信息
+ 要求:无
+ 注意:
  1. 不是in-place操作
  2. 不需要转换的也会复制到新的文件夹下
  3. 如果遇到xml没有对应图片的,则会记录该错误,并生成 ERROR_LOG.txt 文件
"""
from PIL import Image
import os
import xml.etree.ElementTree as ET
import tqdm
import sys


"""============================ 需要修改的地方 ==================================="""
# 输入和输出文件夹路径
XML_PATH  = "EXAMPLE_FOLDER/labels-xml"  # 修正前的 xml 文件夹路径
SAVE_PATH = "EXAMPLE_FOLDER/labels-xml-fixed"  # 修正后的 xml 文件夹路径
IMG_PATH  = "EXAMPLE_FOLDER/images"  # 同名图片文件夹路径
img_type  = '.jpg'  # 图片的格式
"""==============================================================================="""

# 确保输出文件夹存在
if not os.path.exists(SAVE_PATH):
    os.makedirs(SAVE_PATH, exist_ok=True)

# 获取xml文件列表
annotation_files = [file for file in os.listdir(XML_PATH) if file.lower().endswith('.xml')]

"------------计数------------"
TOTAL_NUM   = len(annotation_files)  # 需要处理的 .xml 文件数量
SUCCEED_NUM = 0  # 成功修改的数量
SKIP_NUM    = 0  # 跳过的数量
ERROR_NUM   = 0  # 出错的数量
ERROR_LIST  = []  # 出错的logging
"---------------------------"

# 遍历所有的xml文件
process_bar = tqdm.tqdm(total=TOTAL_NUM, desc="根据图片修正 xml 文件的尺寸 <size> 信息", unit='xml')
for xml_file in annotation_files:
    xml_name, xml_ext = os.path.splitext(xml_file)  # 分离文件名和后缀
    process_bar.set_description(f"Process in \033[1;31m{xml_file}\033[0m")
    
    # 读取 xml 文件
    xml_path = os.path.join(XML_PATH, xml_file)  # 获取完整路径
    tree     = ET.parse(xml_path)  # 解析 xml 树
    root     = tree.getroot()  # 获取 xml 树的根
    
    # 获取同名图片文件名
    image_path = os.path.join(IMG_PATH, xml_name) + img_type
    
    # 判断对应的同名图片文件是否存在,如果不存在则记录错误
    if not os.path.exists(image_path):
        ERROR_NUM += 1
        ERROR_LIST.append(xml_path)
        process_bar.update()
        continue
    
    # 使用PIL获取图片尺寸
    image = Image.open(image_path)
    width, height = image.size
    
    # 判断 xml 中的 <size> 标签是否和图片尺寸对应
    size_elem = root.find("size")
    if size_elem.find("width").text == str(width) and size_elem.find("height").text == str(height):
        # 不需要修正,直接保存文件
        output_path = os.path.join(SAVE_PATH, xml_file)
        tree.write(output_path, encoding="utf-8")
        SKIP_NUM += 1
        process_bar.update()
        continue
    else:
        # 更新xml中的<size>标签
        size_elem.find("width").text = str(width)
        size_elem.find("height").text = str(height)

        # 保存修正后的xml文件
        output_path = os.path.join(SAVE_PATH, xml_file)
        tree.write(output_path, encoding="utf-8")
        SUCCEED_NUM += 1
        process_bar.update()
process_bar.close()

print(f"👌 xml 文件的 size 信息修正已完成, 详情如下:"
      f"\n\t成功修正数量/总xml数量 = \033[1;32m{SUCCEED_NUM}\033[0m/{TOTAL_NUM}"
      f"\n\t跳过数量/总xml数量 = \033[1;34m{SKIP_NUM}\033[0m/{TOTAL_NUM}"
      f"\n\t出错数量/总xml数量 = \033[1;31m{ERROR_NUM}\033[0m/{TOTAL_NUM}")

if SUCCEED_NUM + SKIP_NUM == TOTAL_NUM:
    print("👌 \033[1;32mNo Problems\033[0m")
else:
    print(f"🤡 貌似有点问题, 请仔细核查!"
          f"\n\tSUCCEED_NUM: {SUCCEED_NUM}"
          f"\n\tSKIP_NUM: {SKIP_NUM}"
          f"\n\tERROR_NUM = {ERROR_NUM}"
          f"\nSUCCEED_NUM + SKIP_NUM + ERROR_NUM = {SUCCEED_NUM + SKIP_NUM + ERROR_NUM}"
          f"\nTOTAL_NUM: {TOTAL_NUM}")

if ERROR_LIST:  # 如果有出错信息
    program_path = sys.argv[0]  # 获取程序完整路径
    program_name = os.path.basename(program_path)  # 获取程序名称
    program_parent_path = os.path.dirname(program_path)  # 获取程序所在文件夹路径
    
    ERROR_LOG_PATH = os.path.join(program_parent_path, f"ERROR_LOG-[{program_name}].txt")
    
    with open(ERROR_LOG_PATH, "w") as file:  # 打开文本文件以写入模式
        file.write(f"Program: {program_path}\n")  # 写入程序名称
        file.write(f"🤡 出错了 -> 出错数量/总文件数量 = {ERROR_NUM}/{TOTAL_NUM}\n")  # 写入总体出错信息
        file.write('=' * 50 + '\n')  # 写入分隔线

        # 遍历出错信息列表,写入文件
        for e in ERROR_LIST:
            file.write(f"{e}\n")
            
        # 写入分隔线
        file.write('=' * 50 + '\n')
        
    print(f"\033[1;31m出错信息\033[0m已写入到 [\033[1;34m{ERROR_LOG_PATH}\033[0m] 文件中, 请注意查看!")
相关推荐
AustinCyy3 分钟前
【环境配置】Neo4j Community Windows 安装教程
windows·neo4j
风铃喵游34 分钟前
让大模型调用MCP服务变得超级简单
前端·人工智能
奇怪的杰哥35 分钟前
Win11 加快软件开机自启动
windows
cpsvps1 小时前
Windows内核并发优化
windows
booooooty1 小时前
基于Spring AI Alibaba的多智能体RAG应用
java·人工智能·spring·多智能体·rag·spring ai·ai alibaba
PyAIExplorer1 小时前
基于 OpenCV 的图像 ROI 切割实现
人工智能·opencv·计算机视觉
风口猪炒股指标1 小时前
技术分析、超短线打板模式与情绪周期理论,在市场共识的形成、分歧、瓦解过程中缘起性空的理解
人工智能·博弈论·群体博弈·人生哲学·自我引导觉醒
ai_xiaogui2 小时前
一键部署AI工具!用AIStarter快速安装ComfyUI与Stable Diffusion
人工智能·stable diffusion·部署ai工具·ai应用市场教程·sd快速部署·comfyui一键安装
聚客AI3 小时前
Embedding进化论:从Word2Vec到OpenAI三代模型技术跃迁
人工智能·llm·掘金·日新计划
weixin_387545643 小时前
深入解析 AI Gateway:新一代智能流量控制中枢
人工智能·gateway