76 lines
2.5 KiB
Python
76 lines
2.5 KiB
Python
|
# convert_to_detection.py
|
|||
|
|
import os
|
|||
|
|
import shutil
|
|||
|
|
|
|||
|
|
|
|||
|
|
def convert_labels(input_dir, output_dir=None):
|
|||
|
|
"""
|
|||
|
|
将包含关键点的标签文件转换为仅保留目标检测 bbox 的标签文件
|
|||
|
|
输入格式：class cx cy w h x1 y1 v1 x2 y2 v2 x3 y3 v3 x4 y4 v4
|
|||
|
|
输出格式：class cx cy w h （仅保留前5个）
|
|||
|
|
|
|||
|
|
:param input_dir: 原始标签目录路径，如 'dataset/labels/train'
|
|||
|
|
:param output_dir: 输出目录路径，如果为 None，则覆盖原文件
|
|||
|
|
"""
|
|||
|
|
if output_dir is None:
|
|||
|
|
output_dir = input_dir # 覆盖原文件
|
|||
|
|
print(f"⚠️ 警告：将覆盖原文件目录 {input_dir}")
|
|||
|
|
else:
|
|||
|
|
os.makedirs(output_dir, exist_ok=True)
|
|||
|
|
print(f"✅ 输出目录已创建: {output_dir}")
|
|||
|
|
|
|||
|
|
# 确认输入目录存在
|
|||
|
|
if not os.path.exists(input_dir):
|
|||
|
|
raise FileNotFoundError(f"输入目录不存在: {input_dir}")
|
|||
|
|
|
|||
|
|
# 统计处理文件数
|
|||
|
|
count = 0
|
|||
|
|
|
|||
|
|
for filename in os.listdir(input_dir):
|
|||
|
|
if filename.endswith('.txt'):
|
|||
|
|
input_path = os.path.join(input_dir, filename)
|
|||
|
|
output_path = os.path.join(output_dir, filename)
|
|||
|
|
|
|||
|
|
with open(input_path, 'r') as f:
|
|||
|
|
lines = f.readlines()
|
|||
|
|
|
|||
|
|
new_lines = []
|
|||
|
|
for line in lines:
|
|||
|
|
parts = line.strip().split()
|
|||
|
|
if not parts:
|
|||
|
|
continue
|
|||
|
|
# 只保留前5个数据：class, cx, cy, w, h
|
|||
|
|
det_line = ' '.join(parts[:5])
|
|||
|
|
new_lines.append(det_line)
|
|||
|
|
|
|||
|
|
# 写入新文件
|
|||
|
|
with open(output_path, 'w') as f:
|
|||
|
|
f.write('\n'.join(new_lines))
|
|||
|
|
|
|||
|
|
count += 1
|
|||
|
|
|
|||
|
|
print(f"✅ 转换完成！共处理 {count} 个文件。")
|
|||
|
|
print(f"📁 输出目录: {output_dir}")
|
|||
|
|
|
|||
|
|
|
|||
|
|
# ======================
|
|||
|
|
# 使用示例
|
|||
|
|
# ======================
|
|||
|
|
|
|||
|
|
if __name__ == "__main__":
|
|||
|
|
# ✅ 修改以下路径为你的实际路径
|
|||
|
|
#INPUT_LABELS_DIR = "/media/hx/04e879fa-d697-4b02-ac7e-a4148876ebb0/dataset/point2/train" # 原始标签目录
|
|||
|
|
#OUTPUT_LABELS_DIR = None # 转换后的目录
|
|||
|
|
|
|||
|
|
# 如果你想覆盖原文件，可以设置 OUTPUT_LABELS_DIR = None
|
|||
|
|
# OUTPUT_LABELS_DIR = None
|
|||
|
|
|
|||
|
|
#convert_labels(INPUT_LABELS_DIR, OUTPUT_LABELS_DIR)
|
|||
|
|
|
|||
|
|
# 如果你还有 val 目录，也处理一下
|
|||
|
|
INPUT_VAL_DIR = "/media/hx/04e879fa-d697-4b02-ac7e-a4148876ebb0/dataset/point2/val"
|
|||
|
|
OUTPUT_VAL_DIR = None
|
|||
|
|
|
|||
|
|
convert_labels(INPUT_VAL_DIR, OUTPUT_VAL_DIR)
|
|||
|
|
|
|||
|
|
print("🎉 所有标签已转换完毕！现在可以训练 detect 模型了。")
|