Python divides the training data set and test data set from the Labelme data set by category and proportion

Python divides the training data set and test data set from the Labelme data set according to categories and proportions

  • Preface
  • Prerequisites
  • Related introduction
  • lab environment
  • Split the training data set and test data set from the Labelme data set by category and proportion
    • Code
    • Output results

Foreword

  • Due to my limited level, errors and omissions will inevitably occur. Please criticize and correct me.
  • For more exciting content, you can click to enter the Python Daily Operation Column, OpenCV-Python Small Application Column, YOLO Series Column, Natural Language Processing Column or my personal homepage to view
  • Face disguise detection based on DETR
  • YOLOv7 trains its own data set (mask detection)
  • YOLOv8 trains its own data set (football detection)
  • YOLOv5: TensorRT accelerates YOLOv5 model inference
  • YOLOv5: IoU, GIoU, DIoU, CIoU, EIoU
  • Playing with Jetson Nano (5): TensorRT accelerates YOLOv5 target detection
  • YOLOv5: Add SE, CBAM, CoordAtt, ECA attention mechanism
  • YOLOv5: Interpretation of yolov5s.yaml configuration file and adding small target detection layer
  • Python converts COCO format instance segmentation data set to YOLO format instance segmentation data set
  • YOLOv5: Use version 7.0 to train your own instance segmentation model (instance segmentation of vehicles, pedestrians, road signs, lane lines, etc.)
  • Use Kaggle GPU resources to experience the Stable Diffusion open source project for free

Prerequisite

  • Familiar with Python

Related introduction

  • Python is a cross-platform computer programming language. It is a high-level scripting language that combines interpretation, compilation, interactivity and object-oriented. It was originally designed for writing automated scripts (shells). As the version is constantly updated and new language features are added, it is increasingly used for the development of independent, large-scale projects.
  • PyTorch is a deep learning framework that encapsulates many network and deep learning related tools for our convenience, instead of having to write them one by one. It is divided into CPU and GPU versions. Other frameworks include TensorFlow, Caffe, etc. PyTorch was launched by Facebook Artificial Intelligence Research (FAIR) based on Torch. It is a Python-based sustainable computing package that provides two advanced functions: 1. Tensor computing with powerful GPU acceleration (such as NumPy); 2. , Automatic differentiation mechanism when building deep neural networks.
  • YOLOv5 is a single-stage target detection algorithm. This algorithm adds some new improvement ideas on the basis of YOLOv4, which greatly improves its speed and accuracy. It is a family of object detection architectures and models pre-trained on the COCO dataset and represents Ultralytics’ open-source research into future visual AI approaches, incorporating lessons learned and best practices developed through thousands of hours of research and development. .
  • Labelme is an image annotation tool developed by the Computer Science and Artificial Intelligence Laboratory (CSAIL) of the Massachusetts Institute of Technology (MIT). It is written in Python and PyQT, open source and free. Labelme supports operating systems such as Windows, Linux and Mac.
  • This tool provides an intuitive graphical interface that allows users to annotate many types of objects on images, such as rectangular boxes, polygons, lines, etc., and even more complex shapes. The annotation results are saved in JSON format to facilitate subsequent processing and analysis. This annotation information can be used for tasks such as target detection, image segmentation, and image classification.
  • Overall, Labelme is a powerful and easy-to-use image annotation tool that can meet different image processing needs.
  • Labelme annotation json file is a file format used to store annotation information. It contains the following main fields:
    • version: Labelme’s version number, such as “4.5.6”.
    • flags: Some global flags, such as whether it is a segmentation task, whether it is polygonal, etc.
    • shapes: A list, each element is a dictionary, representing a label object. Each dictionary contains the following fields:
      • label: The category name of the label object, such as “dog”.
      • points: A list, each element is a coordinate pair, representing the boundary points of the annotation object, such as [[10, 20], [30, 40]].
      • group_id: The group number of the labeled object, used to indicate objects belonging to the same group, such as 1.
      • shape_type: The shape type of the annotation object, such as “polygon”, “rectangle”, “circle”, etc.
      • flags: Some flags for the annotation object, such as whether it is a difficult example, whether it is blocked, etc.
    • lineColor: The border line color of the annotation object, such as [0, 255, 0, 128].
    • fillColor: The fill color of the annotation object, such as [255, 0, 0, 128].
    • imagePath: The relative path of the image file, such as “img_001.jpg”.
    • imageData: The binary data of the image file, a string encoded by base64, such as “iVBORw0KGgoAAAANSUhEUgAA…”.
    • imageHeight: The height of the image, for example 600.
    • imageWidth: The width of the image, such as 800.

The following is an example of Labelme labeling a json file:

{<!-- -->
  "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
}

Experimental environment

  • Python 3.x (object-oriented high-level language)

Divide the training data set and test data set from the Labelme data set according to categories and proportions

  • Background: Usually we divide the labeled data set into a training data set and a test data set in a ratio of 8:2.
  • Directory structure example

Code implementation

  • images: The folder where the undivided image data set is located.
  • jsons: The folder where the undivided Labelme annotation files are located.
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 data to variable 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: # Traverse json files
        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 data to variable 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' # The folder where the pictures are located
    in_json_dir = 'jsons' # The folder where json is located

    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)

    # Get all json files
    json_name_list = [i for i in os.listdir(in_json_dir) if i.endswith('.json')] # json file list
    random.shuffle(json_name_list) # shuffle

    all_label_num_dict = {<!-- -->}
    for json_name in json_name_list: # Traverse json files
        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)

Output results

  • images: The folder where the divided training images are located.
  • jsons: The folder where the divided training Labelme annotation files are located.

  • out_images: The folder where the divided test images are located.
  • out_jsons: The folder where the divided test Labelme annotation files are located.
  • Due to my limited level, errors and omissions will inevitably occur. Please criticize and correct me.
  • For more exciting content, you can click to enter the Python Daily Operation Column, OpenCV-Python Small Application Column, YOLO Series Column, Natural Language Processing Column or my personal homepage to view
  • Face disguise detection based on DETR
  • YOLOv7 trains its own data set (mask detection)
  • YOLOv8 trains its own data set (football detection)
  • YOLOv5: TensorRT accelerates YOLOv5 model inference
  • YOLOv5: IoU, GIoU, DIoU, CIoU, EIoU
  • Playing with Jetson Nano (5): TensorRT accelerates YOLOv5 target detection
  • YOLOv5: Add SE, CBAM, CoordAtt, ECA attention mechanism
  • YOLOv5: Interpretation of yolov5s.yaml configuration file and adding small target detection layer
  • Python converts COCO format instance segmentation data set to YOLO format instance segmentation data set
  • YOLOv5: Use version 7.0 to train your own instance segmentation model (instance segmentation of vehicles, pedestrians, road signs, lane lines, etc.)
  • Use Kaggle GPU resources to experience the Stable Diffusion open source project for free