Python将Labelme目标检测数据集和JSON标注文件进行数据增广

在这里插入图片描述

前言

前提条件

相关介绍

  • Python是一种跨平台的计算机程序设计语言。是一个高层次的结合了解释性、编译性、互动性和面向对象的脚本语言。最初被设计用于编写自动化脚本(shell),随着版本的不断更新和语言新功能的添加,越多被用于独立的、大型项目的开发。
  • PyTorch 是一个深度学习框架,封装好了很多网络和深度学习相关的工具方便我们调用,而不用我们一个个去单独写了。它分为 CPU 和 GPU 版本,其他框架还有 TensorFlow、Caffe 等。PyTorch 是由 Facebook 人工智能研究院(FAIR)基于 Torch 推出的,它是一个基于 Python 的可续计算包,提供两个高级功能:1、具有强大的 GPU 加速的张量计算(如 NumPy);2、构建深度神经网络时的自动微分机制。
  • YOLOv5是一种单阶段目标检测算法,该算法在YOLOv4的基础上添加了一些新的改进思路,使其速度与精度都得到了极大的性能提升。它是一个在COCO数据集上预训练的物体检测架构和模型系列,代表了Ultralytics对未来视觉AI方法的开源研究,其中包含了经过数千小时的研究和开发而形成的经验教训和最佳实践。
  • Labelme是一款图像标注工具,由麻省理工(MIT)的计算机科学和人工智能实验室(CSAIL)研发。它是用Python和PyQT编写的,开源且免费。Labelme支持Windows、Linux和Mac等操作系统。
  • 这款工具提供了直观的图形界面,允许用户在图像上标注多种类型的目标,例如矩形框、多边形、线条等,甚至包括更复杂的形状。标注结果以JSON格式保存,便于后续处理和分析。这些标注信息可以用于目标检测、图像分割、图像分类等任务。
  • 总的来说,Labelme是一款强大且易用的图像标注工具,可以满足不同的图像处理需求。
  • Labelme标注json文件是一种用于存储标注信息的文件格式,它包含了以下几个主要的字段:
    • version: Labelme的版本号,例如"4.5.6"。
    • flags: 一些全局的标志,例如是否是分割任务,是否有多边形,等等。
    • shapes: 一个列表,每个元素是一个字典,表示一个标注对象。每个字典包含了以下几个字段:
      • label: 标注对象的类别名称,例如"dog"。
      • points: 一个列表,每个元素是一个坐标对,表示标注对象的边界点,例如[[10, 20], [30, 40]]。
      • group_id: 标注对象的分组编号,用于表示属于同一组的对象,例如1。
      • shape_type: 标注对象的形状类型,例如"polygon",“rectangle”,“circle”,等等。
      • flags: 一些针对该标注对象的标志,例如是否是难例,是否被遮挡,等等。
    • lineColor: 标注对象的边界线颜色,例如[0, 255, 0, 128]。
    • fillColor: 标注对象的填充颜色,例如[255, 0, 0, 128]。
    • imagePath: 图像文件的相对路径,例如"img_001.jpg"。
    • imageData: 图像文件的二进制数据,经过base64编码后的字符串,例如"iVBORw0KGgoAAAANSUhEUgAA…"。
    • imageHeight: 图像的高度,例如600。
    • imageWidth: 图像的宽度,例如800。

以下是一个Labelme标注json文件的示例:

{
  "version": "4.5.6",
  "flags": {},
  "shapes": [
    {
      "label": "dog",
      "points": [
        [
          121.0,
          233.0
        ],
        [
          223.0,
          232.0
        ],
        [
          246.0,
          334.0
        ],
        [
          121.0,
          337.0
        ]
      ],
      "group_id": null,
      "shape_type": "polygon",
      "flags": {}
    }
  ],
  "lineColor": [
    0,
    255,
    0,
    128
  ],
  "fillColor": [
    255,
    0,
    0,
    128
  ],
  "imagePath": "img_001.jpg",
  "imageData": "iVBORw0KGgoAAAANSUhEUgAA...",
  "imageHeight": 600,
  "imageWidth": 800
}

实验环境

  • Python 3.x (面向对象的高级语言)
  • albumentations==1.3.1
  • opencv-python==4.12.0.88

Labelme目标检测数据集和JSON标注文件进行数据增广

  • 背景:将标注好的数据集,随机仿射变换,以达到数据增强的目的。
  • 目录结构示例

在这里插入图片描述

  • png:图片数据集。
  • json:Labelme标注文件。

在这里插入图片描述

{
  "version": "5.5.0",
  "flags": {},
  "shapes": [
    {
      "score": 0.99,
      "label": "2",
      "points": [
        [
          310.379898,
          333.107918
        ],
        [
          429.721062,
          434.07660999999996
        ]
      ],
      "group_id": null,
      "description": null,
      "shape_type": "rectangle",
      "flags": {},
      "mask": null
    },
    {
      "score": 0.99,
      "label": "2",
      "points": [
        [
          68.5886295,
          312.772174
        ],
        [
          229.72838250000004,
          436.130706
        ]
      ],
      "group_id": null,
      "description": null,
      "shape_type": "rectangle",
      "flags": {},
      "mask": null
    },
    {
      "score": 0.99,
      "label": "2",
      "points": [
        [
          208.7640663,
          331.285174
        ],
        [
          250.1050137,
          380.598482
        ]
      ],
      "group_id": null,
      "description": null,
      "shape_type": "rectangle",
      "flags": {},
      "mask": null
    },
    {
      "score": 0.99,
      "label": "7",
      "points": [
        [
          237.97699020000002,
          318.716904
        ],
        [
          286.3156158,
          373.767064
        ]
      ],
      "group_id": null,
      "description": null,
      "shape_type": "rectangle",
      "flags": {},
      "mask": null
    },
    {
      "score": 0.99,
      "label": "2",
      "points": [
        [
          303.30663135000003,
          336.1300624
        ],
        [
          327.16299465,
          362.2645136
        ]
      ],
      "group_id": null,
      "description": null,
      "shape_type": "rectangle",
      "flags": {},
      "mask": null
    },
    {
      "score": 0.99,
      "label": "2",
      "points": [
        [
          284.1491136,
          331.6867246
        ],
        [
          306.15991440000005,
          349.0834754
        ]
      ],
      "group_id": null,
      "description": null,
      "shape_type": "rectangle",
      "flags": {},
      "mask": null
    }
  ],
  "imagePath": "1.png",
  "imageData": null,
  "imageHeight": 484,
  "imageWidth": 579
}

代码实现

import os
import json
import cv2
import numpy as np
import albumentations as A
from pathlib import Path
import random
from tqdm import tqdm

# 支持的图片格式
supported_formats = {'.jpg', '.jpeg', '.png', '.bmp'}

def load_labelme_json(json_path):
    """加载Labelme JSON文件和对应的图像,提取矩形框标注"""
    try:
        with open(json_path, 'r', encoding='utf-8') as f:
            data = json.load(f)
    except json.JSONDecodeError as e:
        print(f"⚠️ JSON解析错误: {json_path} - {str(e)}")
        return None, None, None, None
    
    # 获取图像路径
    image_path = os.path.join(os.path.dirname(json_path), data["imagePath"])
    
    # 检查图像格式
    if not any(image_path.lower().endswith(fmt) for fmt in supported_formats):
        print(f"⚠️ 跳过不支持的图片格式: {image_path}")
        return None, None, None, None
    
    # 加载图像
    image = cv2.imread(image_path)
    if image is None:
        print(f"⚠️ 图片加载失败: {image_path}")
        return None, None, None, None
    
    # 提取矩形框标注
    bboxes = []
    labels = []
    for shape in data["shapes"]:
        # 只处理矩形框
        if shape.get("shape_type", "") != "rectangle":
            continue
            
        # 获取矩形框坐标
        points = np.array(shape["points"], dtype=np.float32)
        x_min, y_min = points[0]
        x_max, y_max = points[1]
        
        # 确保坐标正确性
        x_min, x_max = min(x_min, x_max), max(x_min, x_max)
        y_min, y_max = min(y_min, y_max), max(y_min, y_max)
        
        # 转换为Python原生float类型
        x_min, y_min, x_max, y_max = float(x_min), float(y_min), float(x_max), float(y_max)
        
        # 添加到列表
        bboxes.append([x_min, y_min, x_max, y_max])
        labels.append(shape["label"])
    
    return image, bboxes, labels, data

def get_augmentation_pipeline(transform_type, img_shape):
    """获取针对检测任务的增强管道"""
    height, width = img_shape[:2]
    
    if transform_type == "horizontal_flip":
        return A.Compose([
            A.HorizontalFlip(p=1)
        ], bbox_params=A.BboxParams(format='pascal_voc', label_fields=['class_labels']))
    
    elif transform_type == "vertical_flip":
        return A.Compose([
            A.VerticalFlip(p=1)
        ], bbox_params=A.BboxParams(format='pascal_voc', label_fields=['class_labels']))
    
    elif transform_type == "rotate_90":
        return A.Compose([
            A.RandomRotate90(p=1)
        ], bbox_params=A.BboxParams(format='pascal_voc', label_fields=['class_labels']))
    
    elif transform_type == "rotate":
        # 随机旋转角度 -15到15度
        return A.Compose([
            A.Rotate(limit=15, border_mode=cv2.BORDER_CONSTANT, value=0, p=1)
        ], bbox_params=A.BboxParams(format='pascal_voc', label_fields=['class_labels']))
    
    elif transform_type == "scale":
        # 随机缩放 80%-120%
        return A.Compose([
            A.Affine(scale=(0.8, 1.2), mode=cv2.BORDER_CONSTANT, cval=0, p=1)
        ], bbox_params=A.BboxParams(format='pascal_voc', label_fields=['class_labels']))
    
    elif transform_type == "shift":
        # 随机平移 ±10%
        return A.Compose([
            A.Affine(translate_percent=(-0.1, 0.1), mode=cv2.BORDER_CONSTANT, cval=0, p=1)
        ], bbox_params=A.BboxParams(format='pascal_voc', label_fields=['class_labels']))
    
    elif transform_type == "crop":
        # 随机裁剪 80%-95% 区域
        crop_size = random.uniform(0.8, 0.95)
        return A.Compose([
            A.RandomResizedCrop(
                height=height,
                width=width,
                scale=(crop_size, crop_size),
                ratio=(1.0, 1.0),
                p=1
            )
        ], bbox_params=A.BboxParams(format='pascal_voc', label_fields=['class_labels']))
    
    elif transform_type == "brightness":
        # 随机亮度对比度调整
        return A.Compose([
            A.RandomBrightnessContrast(brightness_limit=0.2, contrast_limit=0.2, p=1)
        ], bbox_params=A.BboxParams(format='pascal_voc', label_fields=['class_labels']))
    
    elif transform_type == "noise":
        # 添加高斯噪声
        return A.Compose([
            A.GaussNoise(var_limit=(10.0, 50.0), p=1)
        ], bbox_params=A.BboxParams(format='pascal_voc', label_fields=['class_labels']))
    
    elif transform_type == "motion_blur":
        # 运动模糊
        return A.Compose([
            A.MotionBlur(blur_limit=7, p=1)
        ], bbox_params=A.BboxParams(format='pascal_voc', label_fields=['class_labels']))
    
    elif transform_type == "color_jitter":
        # 颜色抖动
        return A.Compose([
            A.HueSaturationValue(hue_shift_limit=20, sat_shift_limit=30, val_shift_limit=20, p=1)
        ], bbox_params=A.BboxParams(format='pascal_voc', label_fields=['class_labels']))
    
    else:
        return A.Compose([], bbox_params=A.BboxParams(format='pascal_voc', label_fields=['class_labels']))

def apply_augmentation(transform_type, image, bboxes, labels):
    """应用增强变换到图像和边界框"""
    if not bboxes or not labels:
        return image.copy(), [], []
    
    # 确保bboxes是numpy数组
    bboxes = np.array(bboxes, dtype=np.float32)
    
    # 应用变换
    try:
        transform = get_augmentation_pipeline(transform_type, image.shape)
        transformed = transform(image=image, bboxes=bboxes, class_labels=labels)
    except Exception as e:
        print(f"⚠️ 应用增强时出错 ({transform_type}): {str(e)}")
        return image.copy(), [], []
    
    transformed_image = transformed['image']
    transformed_bboxes = transformed['bboxes']
    transformed_labels = transformed['class_labels']
    
    # 过滤掉无效的边界框
    valid_bboxes = []
    valid_labels = []
    for bbox, label in zip(transformed_bboxes, transformed_labels):
        # 确保坐标是Python原生浮点数
        x_min, y_min, x_max, y_max = bbox
        x_min, y_min, x_max, y_max = float(x_min), float(y_min), float(x_max), float(y_max)
        
        # 检查边界框是否有效
        if x_max - x_min > 2 and y_max - y_min > 2 and 0 <= x_min < x_max and 0 <= y_min < y_max:
            valid_bboxes.append([x_min, y_min, x_max, y_max])
            valid_labels.append(label)
    
    return transformed_image, valid_bboxes, valid_labels

def save_augmented_data(image, bboxes, labels, original_data, output_dir, base_name, index, original_format):
    """保存增强后的图像和标注文件"""
    # 使用原始文件名前缀 + 序号 + 原始后缀
    filename = f"{base_name}_{index}"
    img_path = os.path.join(output_dir, filename + original_format)
    json_path = os.path.join(output_dir, filename + ".json")
    
    # 保存图像
    cv2.imwrite(img_path, image)
    
    # 创建新的JSON数据结构
    new_shapes = []
    for bbox, label in zip(bboxes, labels):
        x_min, y_min, x_max, y_max = bbox
        # 确保所有坐标都是Python原生float类型
        new_shapes.append({
            "label": label,
            "points": [
                [float(x_min), float(y_min)], 
                [float(x_max), float(y_max)]
            ],
            "shape_type": "rectangle",
            "flags": {}
        })
    
    # 确保图像高度宽度是int类型
    height, width = image.shape[:2]
    
    new_data = {
        "version": original_data.get("version", "5.0.1"),
        "flags": original_data.get("flags", {}),
        "shapes": new_shapes,
        "imagePath": filename + original_format,
        "imageData": None,
        "imageHeight": int(height),
        "imageWidth": int(width)
    }
    
    # 保存JSON文件 - 确保所有浮点数都是Python原生类型
    try:
        with open(json_path, 'w', encoding='utf-8') as f:
            json.dump(new_data, f, indent=2, ensure_ascii=False, default=json_float_default)
    except Exception as e:
        print(f"⚠️ 保存JSON失败: {json_path} - {str(e)}")
        # 删除可能已创建的部分文件
        if os.path.exists(img_path):
            os.remove(img_path)
        if os.path.exists(json_path):
            os.remove(json_path)

def json_float_default(obj):
    """自定义JSON序列化函数,处理NumPy类型"""
    if isinstance(obj, np.integer):
        return int(obj)
    elif isinstance(obj, np.floating):
        return float(obj)
    elif isinstance(obj, np.ndarray):
        return obj.tolist()
    else:
        return str(obj)

def detect_augmentation(input_dir, output_dir, iterations_per_image=5):
    """执行检测数据集的数据增强"""
    # 增强类型列表
    transform_types = [
        "horizontal_flip",
        "vertical_flip",
        "rotate_90",
        "rotate",
        "scale",
        "shift",
        "crop",
        "brightness",
        "noise",
        "motion_blur",
        "color_jitter"
    ]
    
    # 确保输出目录存在
    os.makedirs(output_dir, exist_ok=True)
    
    # 获取所有JSON文件
    json_files = [f for f in os.listdir(input_dir) if f.endswith(".json")]
    
    # 处理每个JSON文件
    for json_file in tqdm(json_files, desc="Processing images"):
        json_path = os.path.join(input_dir, json_file)
        
        try:
            # 加载数据和标注
            image, bboxes, labels, data = load_labelme_json(json_path)
            if image is None:
                print(f"⚠️ 图片加载失败: {json_file}")
                continue
            if not bboxes:
                print(f"⚠️ 无有效标注: {json_file}")
            
            # 获取原始图片信息
            image_path = data["imagePath"]
            original_format = Path(image_path).suffix
            if original_format.lower() not in supported_formats:
                print(f"⚠️ 不支持的图片格式: {original_format}")
                continue
            
            # 获取原始文件名前缀
            base_name = Path(image_path).stem
            
            # 保存原始图像作为第一个样本 (序号0)
            save_augmented_data(
                image.copy(), 
                bboxes, 
                labels, 
                data, 
                output_dir, 
                base_name, 
                0, 
                original_format
            )
            
            # 应用增强
            for i in range(1, iterations_per_image + 1):
                # 随机选择增强类型
                trans_type = random.choice(transform_types)
                
                try:
                    # 应用增强
                    augmented_image, augmented_bboxes, augmented_labels = apply_augmentation(
                        trans_type, 
                        image.copy(), 
                        bboxes, 
                        labels
                    )
                    
                    # 检查是否有有效的边界框
                    if not augmented_bboxes:
                        print(f"⚠️ 增强后无有效边界框: {base_name}_{i}{original_format} ({trans_type})")
                        continue
                    
                    # 保存增强结果
                    save_augmented_data(
                        augmented_image, 
                        augmented_bboxes, 
                        augmented_labels, 
                        data, 
                        output_dir, 
                        base_name, 
                        i, 
                        original_format
                    )
                except Exception as e:
                    print(f"⚠️ 增强失败 ({trans_type}): {str(e)}")
        
        except Exception as e:
            print(f"⚠️ 处理文件 {json_file} 时出错: {str(e)}")

    print(f"✅ 数据增强完成! 增强后的数据保存在: {output_dir}")


if __name__ == "__main__":
    # 设置输入输出目录
    input_dir = r"detect"  # 原始数据集路径
    output_dir = r"detect"  # 增强后输出路径
    
    # 控制每张图像增强次数 (不包括原始图像)
    iterations_per_image = 5
    
    # 执行增强
    detect_augmentation(input_dir, output_dir, iterations_per_image)

输出结果

在这里插入图片描述

  • png:图片数据集。
  • json:Labelme标注文件。

在这里插入图片描述
在这里插入图片描述
在这里插入图片描述

评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包

打赏作者

FriendshipT

你的鼓励将是我创作的最大动力

¥1 ¥2 ¥4 ¥6 ¥10 ¥20
扫码支付:¥1
获取中
扫码支付

您的余额不足,请更换扫码支付或充值

打赏作者

实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值