116 lines
4.8 KiB
Python
116 lines
4.8 KiB
Python
import json
|
||
import os
|
||
import glob
|
||
|
||
def labelme_to_yolo_keypoints_batch(json_dir, output_dir, target_label="夹具1", class_id=0, img_shape=None, num_keypoints=None):
|
||
"""
|
||
批量将 LabelMe JSON 文件中的关键点标注转换为 YOLO Pose (Keypoints) 格式的 .txt 文件
|
||
仅转换指定标签（如 "夹具1"），忽略其他标签，且不生成空文件。
|
||
|
||
:param json_dir: 包含 LabelMe JSON 文件的目录
|
||
:param output_dir: 输出 .txt 文件的目录
|
||
:param target_label: 要转换的目标标签名称
|
||
:param class_id: 对应的 YOLO 类别 ID
|
||
:param img_shape: 图像尺寸 (height, width)，如 (1440, 2506)
|
||
:param num_keypoints: 预期关键点数量（可选，用于校验）
|
||
"""
|
||
if img_shape is None:
|
||
raise ValueError("必须提供 img_shape 参数，例如 (1440, 2506)")
|
||
|
||
if num_keypoints is not None and num_keypoints <= 0:
|
||
raise ValueError("num_keypoints 必须为正整数")
|
||
|
||
# 确保输出目录存在
|
||
os.makedirs(output_dir, exist_ok=True)
|
||
|
||
# 获取所有 .json 文件（排除 _mask.json 等）
|
||
json_files = glob.glob(os.path.join(json_dir, "*.json"))
|
||
json_files = [f for f in json_files if os.path.isfile(f) and not f.endswith("_mask.json")]
|
||
|
||
if not json_files:
|
||
print(f"❌ 在 {json_dir} 中未找到任何 JSON 文件")
|
||
return
|
||
|
||
img_h, img_w = img_shape
|
||
converted_count = 0
|
||
skipped_count = 0
|
||
|
||
print(f"🔍 开始转换关键点，仅处理标签: '{target_label}' (class_id={class_id})")
|
||
|
||
for json_file in json_files:
|
||
try:
|
||
with open(json_file, 'r', encoding='utf-8') as f:
|
||
data = json.load(f)
|
||
|
||
base_name = os.path.splitext(os.path.basename(json_file))[0]
|
||
output_path = os.path.join(output_dir, f"{base_name}.txt")
|
||
|
||
has_valid_shapes = False
|
||
keypoints_found = []
|
||
|
||
with open(output_path, 'w', encoding='utf-8') as out_f:
|
||
for shape in data.get('shapes', []):
|
||
label = shape['label']
|
||
points = shape['points'] # 每个 shape 是一个关键点 [x, y]
|
||
|
||
if label != target_label:
|
||
continue
|
||
|
||
# 收集所有关键点坐标
|
||
for x, y in points:
|
||
nx = max(0.0, min(1.0, x / img_w))
|
||
ny = max(0.0, min(1.0, y / img_h))
|
||
keypoints_found.extend([f"{nx:.6f}", f"{ny:.6f}", "2"]) # v=2: visible
|
||
|
||
has_valid_shapes = True
|
||
|
||
# 写入 YOLO 行（支持多个 shape？通常一个标签一个实例）
|
||
if has_valid_shapes:
|
||
if num_keypoints is not None and len(keypoints_found) // 3 != num_keypoints:
|
||
print(f"⚠️ {os.path.basename(json_file)}: 关键点数量不匹配，期望 {num_keypoints}，实际 {len(keypoints_found)//3}")
|
||
line = f"{class_id} {' '.join(keypoints_found)}"
|
||
out_f.write(line + '\n')
|
||
print(f"✅ 已转换: {os.path.basename(json_file)} -> {os.path.basename(output_path)}")
|
||
converted_count += 1
|
||
else:
|
||
# 删除空文件
|
||
os.remove(output_path)
|
||
|
||
if not has_valid_shapes:
|
||
print(f"🟡 跳过: {os.path.basename(json_file)} -> 未包含 '{target_label}'，不生成 .txt")
|
||
skipped_count += 1
|
||
|
||
except Exception as e:
|
||
print(f"❌ 转换失败 {json_file}: {e}")
|
||
if os.path.exists(output_path):
|
||
os.remove(output_path)
|
||
|
||
print("\n" + "="*50)
|
||
print(f"🎉 批量关键点转换完成！")
|
||
print(f"📊 成功转换: {converted_count}")
|
||
print(f"📊 跳过文件: {skipped_count}")
|
||
print(f"📁 输出目录: {output_dir}")
|
||
if num_keypoints:
|
||
print(f"📍 关键点数量: {num_keypoints}")
|
||
print("="*50)
|
||
|
||
|
||
# ================== 用户配置区 ==================
|
||
JSON_DIR = "/home/hx/yolo/yolo11_point/folder_end" # LabelMe JSON 文件夹路径
|
||
OUTPUT_DIR = "labels_keypoints" # 输出 YOLO 关键点标签目录
|
||
TARGET_LABEL = "point" # 要提取的关键点标签
|
||
CLASS_ID = 0 # YOLO 中该类的 ID
|
||
IMG_SHAPE = (1440, 2506) # (高度, 宽度)
|
||
NUM_KEYPOINTS = 4 # 可选：指定关键点数量用于校验
|
||
|
||
# ================== 执行转换 ==================
|
||
if __name__ == "__main__":
|
||
print(f"🚀 开始转换 LabelMe 关键点 → YOLO Pose 格式 (仅 '{TARGET_LABEL}')")
|
||
labelme_to_yolo_keypoints_batch(
|
||
json_dir=JSON_DIR,
|
||
output_dir=OUTPUT_DIR,
|
||
target_label=TARGET_LABEL,
|
||
class_id=CLASS_ID,
|
||
img_shape=IMG_SHAPE,
|
||
num_keypoints=NUM_KEYPOINTS
|
||
) |