lableme标注的图片转成灰度的图片用于数据扩增

点击查看代码

import os
import json
import cv2
import numpy as np
from labelme import utils


def process_single_json(json_path, output_dir):
    """
    处理单个 JSON：
    1. 读取原图 -> 转灰度 -> 转回3通道(保持兼容)
    2. 保存新图
    3. 生成指向新图的新 JSON
    """
    # 1. 读取原始 JSON
    with open(json_path, 'r', encoding='utf-8') as f:
        data = json.load(f)

    # 2. 获取原图
    img = None
    # 优先尝试从 JSON 内部读取 imageData
    if data.get('imageData'):
        try:
            img = utils.img_b64_to_arr(data['imageData'])
        except:
            pass

    # 如果内部没有，则从文件读取
    if img is None:
        json_dir = os.path.dirname(json_path)
        img_path = os.path.join(json_dir, data['imagePath'])
        if os.path.exists(img_path):
            img = cv2.imread(img_path)  # 读取为 BGR
        else:
            print(f"❌ 找不到图片: {img_path}")
            return False

    # 3. 图像处理流程: 原图 -> 灰度 -> 3通道灰度图
    # 如果是彩色图 (H, W, 3)，先转灰度
    if len(img.shape) == 3:
        gray_img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    else:
        gray_img = img  # 已经是灰度图

    # 【关键步骤】将单通道灰度图转为 3 通道
    # 这样生成的图片形状是 (H, W, 3)，和原图一致，防止训练报错
    gray_img_3ch = cv2.cvtColor(gray_img, cv2.COLOR_GRAY2BGR)

    # 4. 保存新图片
    os.makedirs(output_dir, exist_ok=True)
    json_filename = os.path.basename(json_path).replace('.json', '')

    # 新图片命名：原文件名_gray.png
    new_img_name = f"{json_filename}_gray.jpg"
    new_img_path = os.path.join(output_dir, new_img_name)

    cv2.imwrite(new_img_path, gray_img_3ch)

    # 5. 更新 JSON
    # 修改 imagePath
    data['imagePath'] = new_img_name

    # 【关键步骤】清空 imageData
    # 设为 None 强制 Labelme 下次打开时去读取硬盘上的新图片，
    # 避免因为旧 JSON 里存的 Base64 数据（原图）和新图片不一致导致显示错误
    data['imageData'] = None

    # 6. 保存新 JSON
    new_json_name = f"{json_filename}_gray.json"
    new_json_path = os.path.join(output_dir, new_json_name)

    with open(new_json_path, 'w', encoding='utf-8') as f:
        json.dump(data, f, ensure_ascii=False, indent=2)

    return True


def batch_run(input_dir, output_dir):
    print(f"📂 正在扫描目录: {input_dir}")
    success_count = 0

    for root, dirs, files in os.walk(input_dir):
        # 保持目录结构
        rel_path = os.path.relpath(root, input_dir)
        target_dir = os.path.join(output_dir, rel_path)
        os.makedirs(target_dir, exist_ok=True)

        for file in files:
            if file.endswith('.json'):
                json_path = os.path.join(root, file)
                if process_single_json(json_path, target_dir):
                    success_count += 1
                    print(f"✅ 处理: {file}")
                else:
                    print(f"❌ 失败: {file}")

    print(f"\n🎉 完成！共处理 {success_count} 个文件。")
    print(f"📁 输出位置: {output_dir}")


# ================= 运行配置 =================
if __name__ == '__main__':
    # 输入目录（包含原图和JSON的文件夹）
    INPUT_DIR = r'D:\pic\Guangzhoumeiwei\num17youmoHeight(double pad)\dataset20260319\src'
    # 输出目录（生成的灰度图数据集）
    OUTPUT_DIR = r'D:\pic\Guangzhoumeiwei\num17youmoHeight(double pad)\dataset20260319\gray'

    batch_run(INPUT_DIR, OUTPUT_DIR)