Files
zjsh_yolov11/ultralytics_yolo11-main/dataset/trans.py
2025-08-13 12:53:33 +08:00

82 lines
3.0 KiB
Python

import json
import os
import glob
def labelme_to_yolo_segmentation_batch(json_dir, output_dir, class_mapping, img_shape):
"""
批量将 LabelMe JSON 文件转换为 YOLO 分割格式的 .txt 文件
:param json_dir: 包含 JSON 文件的目录
:param output_dir: 输出 .txt 文件的目录
:param class_mapping: 类别映射字典,如 {"夹具": 0, "夹具1": 1}
:param img_shape: 图像尺寸 (height, width)
"""
# 确保输出目录存在
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)]
if not json_files:
print(f"❌ 在 {json_dir} 中未找到任何 JSON 文件")
return
img_h, img_w = img_shape
converted_count = 0
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")
with open(output_path, 'w', encoding='utf-8') as out_f:
for shape in data['shapes']:
label = shape['label']
points = shape['points']
if label not in class_mapping:
print(f"⚠️ 跳过未知标签 '{label}' in {json_file}")
continue
class_id = class_mapping[label]
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}")
line = f"{class_id} {' '.join(normalized)}"
out_f.write(line + '\n')
print(f"✅ 已转换: {os.path.basename(json_file)} -> {os.path.basename(output_path)}")
converted_count += 1
except Exception as e:
print(f"❌ 转换失败 {json_file}: {e}")
print(f"\n🎉 批量转换完成!共处理 {converted_count} 个文件。")
print(f"📁 输出目录: {output_dir}")
# ================== 用户配置区 ==================
JSON_DIR = "/home/hx/yolo/ultralytics_yolo11-main/dataset" # 输入:存放 LabelMe JSON 的文件夹
OUTPUT_DIR = "labels" # 输出:存放 YOLO .txt 的文件夹
CLASS_MAPPING = {
"夹具": 0,
"夹具1": 1
# 可继续添加其他类别
}
IMG_SHAPE = (1440, 2506) # 替换为你的图像实际尺寸 (height, width)
# ================== 执行转换 ==================
if __name__ == "__main__":
labelme_to_yolo_segmentation_batch(
json_dir=JSON_DIR,
output_dir=OUTPUT_DIR,
class_mapping=CLASS_MAPPING,
img_shape=IMG_SHAPE
)