109 lines
4.4 KiB
Python
109 lines
4.4 KiB
Python
import json
|
||
import os
|
||
import glob
|
||
|
||
def labelme_to_yolo_segmentation_batch(json_dir, output_dir, target_label="夹具1", class_id=1, img_shape=None):
|
||
"""
|
||
批量将 LabelMe JSON 文件转换为 YOLO 分割格式的 .txt 文件
|
||
仅转换指定标签（如 "夹具1"），忽略其他所有标签，且不生成空文件。
|
||
|
||
:param json_dir: 包含 LabelMe JSON 文件的目录
|
||
:param output_dir: 输出 .txt 文件的目录
|
||
:param target_label: 要转换的目标标签名称，如 "夹具1"
|
||
:param class_id: 对应的 YOLO 类别 ID（通常从 0 开始）
|
||
:param img_shape: 图像尺寸 (height, width)，如 (1440, 2506)
|
||
"""
|
||
if img_shape is None:
|
||
raise ValueError("必须提供 img_shape 参数，例如 (1440, 2506)")
|
||
|
||
# 确保输出目录存在
|
||
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 # 标记是否有目标标签
|
||
|
||
# 打开输出文件准备写入
|
||
with open(output_path, 'w', encoding='utf-8') as out_f:
|
||
for shape in data.get('shapes', []):
|
||
label = shape['label']
|
||
points = shape['points']
|
||
|
||
# 只处理目标标签
|
||
if label != target_label:
|
||
continue
|
||
|
||
# 归一化坐标
|
||
normalized = []
|
||
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))
|
||
normalized.append(f"{nx:.6f}")
|
||
normalized.append(f"{ny:.6f}")
|
||
|
||
# 写入 YOLO 行
|
||
line = f"{class_id} {' '.join(normalized)}"
|
||
out_f.write(line + '\n')
|
||
has_valid_shapes = True
|
||
|
||
# 如果没有找到目标标签，删除空文件
|
||
if not has_valid_shapes:
|
||
os.remove(output_path)
|
||
print(f"🟡 跳过: {os.path.basename(json_file)} -> 未包含 '{target_label}'，不生成 .txt")
|
||
skipped_count += 1
|
||
else:
|
||
print(f"✅ 已转换: {os.path.basename(json_file)} -> {os.path.basename(output_path)}")
|
||
converted_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}")
|
||
print("="*50)
|
||
|
||
|
||
# ================== 用户配置区 ==================
|
||
JSON_DIR = "/media/hx/04e879fa-d697-4b02-ac7e-a4148876ebb0/dataset/seg/class4" # LabelMe JSON 文件夹路径
|
||
OUTPUT_DIR = "/media/hx/04e879fa-d697-4b02-ac7e-a4148876ebb0/dataset/seg/label" # 输出 YOLO 标注文件夹
|
||
TARGET_LABEL = "yeimian" # 只转换这个标签
|
||
#TARGET_LABEL = "夹具1" # 只转换这个标签i
|
||
CLASS_ID = 0 # YOLO 中该类的 ID
|
||
IMG_SHAPE = (1440, 2560) # 图像实际尺寸 (高度, 宽度)
|
||
|
||
# ================== 执行转换 ==================
|
||
if __name__ == "__main__":
|
||
print(f"🚀 开始转换 LabelMe → YOLO (仅 '{TARGET_LABEL}')")
|
||
labelme_to_yolo_segmentation_batch(
|
||
json_dir=JSON_DIR,
|
||
output_dir=OUTPUT_DIR,
|
||
target_label=TARGET_LABEL,
|
||
class_id=CLASS_ID,
|
||
img_shape=IMG_SHAPE
|
||
) |