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 模型了。") |