目标检测标注格式转换：labelme2yolo & yolo2labelme

文章目录[隐藏]

labelme 标注的结果转为 yolo
yolo TTA结果转为 labeme 格式用于人工微调，扩大数据规模

《labelme2yolo & yolo2labelme》

训练 yolo 系列模型的时候需要用到 labelme 标注一些数据，或者用众包标注大量数据。有些时候目标检测的场景比较简单的时候完全可以手工标注一些数据训练一个模型，然后用 TTA 扩充数据量，这就需要我们用到数据标签的格式转换的脚本，这里提供两个 labelme 结果到 yolo 格式的脚本，以及 yolo 结果到 labelme 格式的脚本。

Key Words：labelme、yolo、TTA

Beijing, 2021

作者：mxsurui

Agile Pioneer

labelme 标注的结果转为 yolo

# coding:utf-8

import os
import cv2
import json
import matplotlib.pyplot as plt

"""
1. One row per object
2. Each row is class x_center y_center width height format.
3. Box coordinates must be in normalized xywh format (from 0 - 1). 
If your boxes are in pixels, divide x_center and width by image width, and y_center and height by image height.
4. Class numbers are zero-indexed (start from 0).
"""

# labelme 中预设的类别名和类别 id 的对应关系
label_idx_map = {"0": 0, "1": 1, "2": 2, "3": 3, "4": 4}
color_list = [[200, 0, 0], [0, 200, 0], [0, 0, 200], [200, 200, 0], [0, 200, 200], [200, 0, 200], [0, 0, 0],
              [128, 128, 0]]


def labelme_to_yolo(img_dir, json_dir, save_dir):
    name_list = os.listdir(json_dir)
    for name in name_list:
        if name.startswith('.'):
            continue
        save_path = os.path.join(save_dir, name.replace(".json", ".txt"))
        im_path = os.path.join(img_dir, name.replace(".json", ".png"))
        json_path = os.path.join(json_dir, name)
        im = cv2.imread(im_path)
        label_dict = json.load(open(json_path, 'r'))
        height = label_dict["imageHeight"]
        width = label_dict["imageWidth"]
        loc_info_list = label_dict["shapes"]
        label_info_list = list()
        for loc_info in loc_info_list:
            obj_name = loc_info.get("label")
            label_id = label_idx_map.get(obj_name)
            # print(label_id)
            loc = loc_info.get("points")
            x0, y0 = loc[0]  # 左上角点
            x1, y1 = loc[1]  # 右下角点
            cv2.rectangle(im, (int(x0), int(y0)), (int(x1), int(y1)), color_list[label_id], 2)
            x_center = (x0 + x1) / 2 / width
            y_center = (y0 + y1) / 2 / height
            box_w = (abs(x1 - x0)) / width  # 这里使用绝对值是因为有时候先标注的右下角点
            box_h = (abs(y1 - y0)) / height
            assert box_w > 0, print((int(x0), int(y0)), (int(x1), int(y1)))
            assert box_h > 0
            label_info_list.append([str(label_id), str(x_center), str(y_center), str(box_w), str(box_h)])

        with open(save_path, 'a') as f:
            for label_info in label_info_list:
                label_str = ' '.join(label_info)
                f.write(label_str)
                f.write('\n')

        # debug
        # plt.figure(0)
        # plt.imshow(im)
        # plt.show()
        # print("xxx")


if __name__ == "__main__":
	# 图像文件夹
    image_dir = "/Users/tianshu/Documents/iCollections/tta_train/img"
    # labelme 的标注结果
    json_dir = "/Users/tianshu/Documents/iCollections/tta_train/json"
    # yolo 使用的 txt 结果
    save_dir = "/Users/tianshu/Documents/iCollections/tta_train/txt"

    labelme_to_yolo(image_dir, json_dir, save_dir)

yolo TTA结果转为 labeme 格式用于人工微调，扩大数据规模

# coding:utf-8

import os
import cv2
import json
import matplotlib.pyplot as plt
import base64

img_dir = "/Users/tianshu/Documents/iCollections/tta_train/img"


def parse_tta_label(txt_path, img_dir, save_dir):
    file_name = txt_path.split('/')[-1].split('.')[0]
    img_path = os.path.join(img_dir, file_name + ".png")
    img = cv2.imread(img_path)
    h, w = img.shape[:2]

    with open(img_path, 'rb') as f:
        image = f.read()
    image_base64 = str(base64.b64encode(image), encoding='utf-8')

    with open(txt_path, 'r') as f:
        label_info_list = f.readlines()

    version = "4.5.13"
    data_dict = dict()
    data_dict.__setitem__("version", version)
    data_dict.__setitem__("imagePath", file_name)
    data_dict.__setitem__("imageData", image_base64)
    data_dict.__setitem__("imageHeight", h)
    data_dict.__setitem__("imageWidth", w)
    data_dict.__setitem__("flags", {})
    data_dict["shapes"] = list()
    for label_info in label_info_list:
        label_info = label_info.strip()
        label_info = label_info.split(' ')
        class_name = label_info[0]
        c_x = int(float(label_info[1]) * w)
        c_y = int(float(label_info[2]) * h)
        b_w = int(float(label_info[3]) * w)
        b_h = int(float(label_info[4]) * h)
        x1 = int(c_x - b_w / 2)
        x2 = int(c_x + b_w / 2)
        y1 = int(c_y - b_h / 2)
        y2 = int(c_y + b_h / 2)

        points = [[x1, y1], [x2, y2]]
        shape_type = "rectangle"
        shape = {}

        shape.__setitem__("label", class_name)
        shape.__setitem__("points", points)
        shape.__setitem__("shape_type", shape_type)
        shape.__setitem__("flags", {})
        shape.__setitem__("group_id", None)
        data_dict["shapes"].append(shape)

    save_json_path = os.path.join(save_dir, "%s.json" % file_name)
    json.dump(data_dict, open(save_json_path, 'w'), indent=4)
    #### for view debug
    # img = cv2.rectangle(img, (x1, y1), (x2, y2), (0, 255, 0), 2)
    # plt.figure(0)
    # plt.imshow(img)
    # plt.show()


def generate_labelme_prelabel(txt_dir, img_dir, save_dir):
    txt_name_list = os.listdir(txt_dir)
    for txt_name in txt_name_list:
        if txt_name.startswith('.'):
            continue
        print("processing -> %s" % txt_name)
        txt_path = os.path.join(txt_dir, txt_name)
        parse_tta_label(txt_path, img_dir, save_dir)


if __name__ == '__main__':
    txt_dir = "/Users/tianshu/Documents/iCollections/tta_train/txt"
    save_dir = "/Users/tianshu/Documents/iCollections/tta_train/json"
    generate_labelme_prelabel(txt_dir, img_dir, save_dir)