Python按类别和比例从Labelme数据集中划分出训练数据集和测试数据集

前言

• 由于本人水平有限，难免出现错漏，敬请批评改正。
• 更多精彩内容，可点击进入Python日常小操作专栏、OpenCV-Python小应用专栏、YOLO系列专栏、自然语言处理专栏或我的个人主页查看
• 基于DETR的人脸伪装检测
• YOLOv7训练自己的数据集（口罩检测）
• YOLOv8训练自己的数据集（足球检测）
• YOLOv5：TensorRT加速YOLOv5模型推理
• YOLOv5：IoU、GIoU、DIoU、CIoU、EIoU
• 玩转Jetson Nano（五）：TensorRT加速YOLOv5目标检测
• YOLOv5：添加SE、CBAM、CoordAtt、ECA注意力机制
• YOLOv5：yolov5s.yaml配置文件解读、增加小目标检测层
• Python将COCO格式实例分割数据集转换为YOLO格式实例分割数据集
• YOLOv5：使用7.0版本训练自己的实例分割模型（车辆、行人、路标、车道线等实例分割）
• 使用Kaggle GPU资源免费体验Stable Diffusion开源项目

前提条件

• 熟悉Python

相关介绍

• 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 （面向对象的高级语言）

按类别和比例从Labelme数据集中划分出训练数据集和测试数据集

• 背景：通常我们会将标注好的数据集，按8：2的比列划分为训练数据集和测试数据集。
• 目录结构示例
  

代码实现

• images：未划分的图片数据集所在的文件夹。
• jsons：未划分的Labelme标注文件所在的文件夹。

import os
import cv2
import json
import time
import math
import shutil
import random


def count_json_label(in_json_path,all_label_num_dict):
    with open(in_json_path, "r", encoding='utf-8') as f:
        # json.load数据到变量json_data
        json_data = json.load(f)
    for i in json_data['shapes']:
        if i['label'] in all_label_num_dict.keys():
            all_label_num_dict[i['label']] = all_label_num_dict[i['label']] + 1
        else:
            all_label_num_dict[i['label']] = 1
    all_label_num_dict = dict(sorted(all_label_num_dict.items(), key=lambda x: x[1]))
    return all_label_num_dict


def cal_test_count_json_label(all_label_num_dict,split_rate=0.2):
    test_label_num_dict = {key: math.floor(value * split_rate) for key, value in all_label_num_dict.items()} 
    return test_label_num_dict


def select_test_img_json(in_img_dir,in_json_dir,json_name_list,key,value):
    n = 0
    for json_name in json_name_list: # 遍历json文件
        in_img_path = os.path.join(in_img_dir,json_name[:-5]+img_type)
        in_json_path = os.path.join(in_json_dir,json_name)
        try:
            with open(in_json_path, "r", encoding='utf-8') as f:
                # json.load数据到变量json_data
                json_data = json.load(f)
            if n == value:
                break
            for i in json_data['shapes']:
                if i['label']==key:
                    # print(in_json_path)
                    shutil.move(in_img_path,out_img_dir)
                    shutil.move(in_json_path,out_json_dir)
                    n +=1
                    break
        except Exception as e:
            continue

if __name__=="__main__":
    img_type = '.jpg'
    in_img_dir = 'images' # 图片所在文件夹
    in_json_dir = 'jsons' # json所在文件夹

    out_img_dir = 'out_images'
    if not os.path.exists(out_img_dir):
        os.mkdir(out_img_dir)
    
    out_json_dir = 'out_jsons'
    if not os.path.exists(out_json_dir):
        os.mkdir(out_json_dir)

    # 获取所有json文件
    json_name_list = [i for i in os.listdir(in_json_dir) if i.endswith('.json')] # json文件列表
    random.shuffle(json_name_list) # 乱序  

    all_label_num_dict = {}
    for json_name in json_name_list: # 遍历json文件
        in_json_path = os.path.join(in_json_dir,json_name)
        count_json_label(in_json_path,all_label_num_dict)

    test_label_num_dict = cal_test_count_json_label(all_label_num_dict,split_rate=0.2)

    for key,value in test_label_num_dict.items():
        select_test_img_json(in_img_dir,in_json_dir,json_name_list,key,value)

输出结果

• images：划分出的训练图片所在的文件夹。
• jsons：划分出的训练Labelme标注文件所在的文件夹。

• out_images：划分出的测试图片所在的文件夹。
• out_jsons：划分出的测试Labelme标注文件所在的文件夹。

• 由于本人水平有限，难免出现错漏，敬请批评改正。
• 更多精彩内容，可点击进入Python日常小操作专栏、OpenCV-Python小应用专栏、YOLO系列专栏、自然语言处理专栏或我的个人主页查看
• 基于DETR的人脸伪装检测
• YOLOv7训练自己的数据集（口罩检测）
• YOLOv8训练自己的数据集（足球检测）
• YOLOv5：TensorRT加速YOLOv5模型推理
• YOLOv5：IoU、GIoU、DIoU、CIoU、EIoU
• 玩转Jetson Nano（五）：TensorRT加速YOLOv5目标检测
• YOLOv5：添加SE、CBAM、CoordAtt、ECA注意力机制
• YOLOv5：yolov5s.yaml配置文件解读、增加小目标检测层
• Python将COCO格式实例分割数据集转换为YOLO格式实例分割数据集
• YOLOv5：使用7.0版本训练自己的实例分割模型（车辆、行人、路标、车道线等实例分割）
• 使用Kaggle GPU资源免费体验Stable Diffusion开源项目