41 lines
2.7 KiB
Python
41 lines
2.7 KiB
Python
import os
|
|
|
|
|
|
# 类别名称列表
|
|
class_names = ['badaoxie', 'banxie', 'baomihuaxie', 'heiyaoshixie', 'huoxingxie', 'jianzhenxie', 'maozhuaxie', 'mianhuatangxie', 'mofangxie', 'moshutiexie', 'qidianxie', 'ranshaoxie', 'shuinixie', 'wangmianxie', 'xiaobaixie', 'xiaochujuxie', 'xiaomaixie', 'yingtangxie', 'yinsuxie', 'yuanyangxie', 'yucixie'] # 替换为你的类别名称
|
|
target_class = ['badaoxie', 'banxie', 'baomihuaxie', 'heiyaoshixie', 'huoxingxie', 'jianzhenxie', 'maozhuaxie', 'mianhuatangxie', 'mofangxie', 'moshutiexie', 'qidianxie', 'ranshaoxie', 'shuinixie', 'wangmianxie', 'xiaobaixie', 'xiaochujuxie', 'xiaomaixie', 'yingtangxie', 'yinsuxie', 'yuanyangxie', 'yucixie'] # 目标类别
|
|
target_class_id = {'badaoxie': 2, 'banxie':2, 'baomihuaxie':2, 'heiyaoshixie':2, 'huoxingxie':2,
|
|
'jianzhenxie':2, 'maozhuaxie':2, 'mianhuatangxie':2, 'mofangxie':2, 'moshutiexie':2,
|
|
'qidianxie':2, 'ranshaoxie':2, 'shuinixie':2, 'wangmianxie':2, 'xiaobaixie':2,
|
|
'xiaochujuxie':2, 'xiaomaixie':2, 'yingtangxie':2, 'yinsuxie':2, 'yuanyangxie':2, 'yucixie':2} # 目标类别的新 ID
|
|
|
|
def filter_and_rename_labels(label_directory, output_label_directory):
|
|
if not os.path.exists(output_label_directory):
|
|
os.makedirs(output_label_directory) # 创建输出标签目录
|
|
|
|
# 遍历标签目录中的所有文件
|
|
for label_file in os.listdir(label_directory):
|
|
if label_file.endswith('.txt'): # 确保是标签文件
|
|
with open(os.path.join(label_directory, label_file), 'r') as file:
|
|
lines = file.readlines()
|
|
|
|
# 过滤标签，只保留目标类别并修改类别 ID
|
|
filtered_lines = []
|
|
for line in lines:
|
|
try:
|
|
class_id, x_center, y_center, box_width, box_height = map(float, line.split())
|
|
if class_names[int(class_id)] in target_class:
|
|
# 修改类别 ID 为 0
|
|
filtered_lines.append(f"{target_class_id[class_names[int(class_id)]]} {x_center} {y_center} {box_width} {box_height}\n")
|
|
except:
|
|
print("标签文件错误", label_file)
|
|
|
|
# 将过滤后的标签写入新的文件
|
|
with open(os.path.join(output_label_directory, label_file), 'w') as file:
|
|
file.writelines(filtered_lines)
|
|
|
|
# 示例用法
|
|
label_directory = '/home/admin-root/haotian/python哈汽锻8安全帽识别/Shoes/shoe.v2i.yolov8/valid/labels_raw' # 替换为你的标签文件夹路径
|
|
output_label_directory = '/home/admin-root/haotian/python哈汽锻8安全帽识别/Shoes/shoe.v2i.yolov8/valid/labels' # 替换为你希望保存过滤后标签的路径
|
|
filter_and_rename_labels(label_directory, output_label_directory)
|