【LogoDetection 数据集处理】（5）数据集格式转换：labelme格式转为coco格式

关于labelme与coco格式的介绍网上有很多文章，这里就不做介绍了。可以参考 coco标注信息与labelme标注信息的详解、相互转换及可视化

labelme格式数据集是一张图片对应一个json文件：

coco格式的数据集是将所有图片的信息都放在一个json文件中。

我们要处理的原始数据格式如下：

一个子文件夹代表一个类别。

最终得到的coco格式的图片格式如下，即所有类别的图片都放在一个文件夹中：

要将所有图片放到一个文件夹中，并且为了防止名字重复，原来的 0001.jpg这样的文件必须要重命名为类似于 3ce_0001.jpg，这样可以保证图像文件名是唯一的。

所以转换完之后得到的是一个包含所有图片的文件夹和一个包含所有图片标注信息的json文件。

代码如下：

"""
将labelme格式转为coco格式，并在图中画出GT框。
"""""

import json
import os
import shutil
import argparse
import cv2
import numpy as np
from tqdm import tqdm

_COLORS = np.array(
    [
        0.000, 0.447, 0.741,
        0.850, 0.325, 0.098,
        0.929, 0.694, 0.125,
        0.494, 0.184, 0.556,
        0.466, 0.674, 0.188,
        0.301, 0.745, 0.933,
        0.635, 0.078, 0.184,
        0.300, 0.300, 0.300,
        0.600, 0.600, 0.600,
        1.000, 0.000, 0.000,
        1.000, 0.500, 0.000,
        0.749, 0.749, 0.000,
        0.000, 1.000, 0.000,
        0.000, 0.000, 1.000,
        0.667, 0.000, 1.000,
        0.333, 0.333, 0.000,
        0.333, 0.667, 0.000,
        0.333, 1.000, 0.000,
        0.667, 0.333, 0.000,
        0.667, 0.667, 0.000,
        0.667, 1.000, 0.000,
        1.000, 0.333, 0.000,
        1.000, 0.667, 0.000,
        1.000, 1.000, 0.000,
        0.000, 0.333, 0.500,
        0.000, 0.667, 0.500,
        0.000, 1.000, 0.500,
        0.333, 0.000, 0.500,
        0.333, 0.333, 0.500,
        0.333, 0.667, 0.500,
        0.333, 1.000, 0.500,
        0.667, 0.000, 0.500,
        0.667, 0.333, 0.500,
        0.667, 0.667, 0.500,
        0.667, 1.000, 0.500,
        1.000, 0.000, 0.500,
        1.000, 0.333, 0.500,
        1.000, 0.667, 0.500,
        1.000, 1.000, 0.500,
        0.000, 0.333, 1.000,
        0.000, 0.667, 1.000,
        0.000, 1.000, 1.000,
        0.333, 0.000, 1.000,
        0.333, 0.333, 1.000,
        0.333, 0.667, 1.000,
        0.333, 1.000, 1.000,
        0.667, 0.000, 1.000,
        0.667, 0.333, 1.000,
        0.667, 0.667, 1.000,
        0.667, 1.000, 1.000,
        1.000, 0.000, 1.000,
        1.000, 0.333, 1.000,
        1.000, 0.667, 1.000,
        0.333, 0.000, 0.000,
        0.500, 0.000, 0.000,
        0.667, 0.000, 0.000,
        0.833, 0.000, 0.000,
        1.000, 0.000, 0.000,
        0.000, 0.167, 0.000,
        0.000, 0.333, 0.000,
        0.000, 0.500, 0.000,
        0.000, 0.667, 0.000,
        0.000, 0.833, 0.000,
        0.000, 1.000, 0.000,
        0.000, 0.000, 0.167,
        0.000, 0.000, 0.333,
        0.000, 0.000, 0.500,
        0.000, 0.000, 0.667,
        0.000, 0.000, 0.833,
        0.000, 0.000, 1.000,
        0.000, 0.000, 0.000,
        0.143, 0.143, 0.143,
        0.286, 0.286, 0.286,
        0.429, 0.429, 0.429,
        0.571, 0.571, 0.571,
        0.714, 0.714, 0.714,
        0.857, 0.857, 0.857,
        0.000, 0.447, 0.741,
        0.314, 0.717, 0.741,
        0.50, 0.5, 0
    ]
).astype(np.float32).reshape(-1, 3)

parser = argparse.ArgumentParser(description='labelme_2_coco.py')
parser.add_argument('--rename_file', action='store_true')
parser.add_argument('--visual', action='store_true')
args = parser.parse_args()

root_path = "logo_det0526"  # 原数据文件

# 对原来的所有图片重命名，新的文件名为：子文件夹名+"_"+原文件名
# 该操作只需要第一次运行程序时执行
if args.rename_file:
    subdir_list = os.listdir(root_path)  # root_path下的所有文件名（包括文件和文件夹）
    for subdir_name in tqdm(subdir_list, desc="rename pic and json files"):
        subdirname_length = len(subdir_name)
        subdir_path = os.path.join(root_path, subdir_name)
        # 判断是否是文件夹
        if os.path.isdir(subdir_path):
            file_list = os.listdir(subdir_path)  # subdir_path文件夹下的所有jpg和json文件
            for file_name in file_list:
                # 判断是否需要重命名，如果原来文件名就是3ce_0000.jpg，则不需要重命名
                # 如果原来文件名就是0000.jpg，则需要重命名为3ce_0000.jpg
                if file_name[:subdirname_length] == subdir_name:
                    pass
                else:
                    file_path = os.path.join(subdir_path, file_name)
                    new_filename = subdir_name + '_' + file_name
                    new_filename_path = os.path.join(subdir_path, new_filename)
                    os.rename(file_path, new_filename_path)

# 存放全部图片的新文件夹
new_root_path = 'logo_det0526_new'
isExists = os.path.exists(new_root_path)
if isExists:
    shutil.rmtree(new_root_path)
os.mkdir(new_root_path)

# 创建类别id
categories_list = []
subdir_list = os.listdir(root_path)
categories_num = 1
for subdir_name in subdir_list:
    categories_dict = {}
    categories_dict["id"] = categories_num
    categories_dict["name"] = subdir_name
    categories_list.append(categories_dict)
    categories_num += 1

# coco格式的json文件中的字典
train_dict = {"info": "logo_det0526",
              "license": ["license"],
              "images": [],
              "annotations": [],
              "categories": []
              }
train_dict["categories"] = categories_list

image_ext = ["jpg", "jpeg", "webp", "bmp", "png"]
json_ext = 'json'

train_dict_images_id = 1
train_dict_annotations_id = 1
for subdir_name in tqdm(subdir_list, desc="convert into coco json"):
    subdir_path = os.path.join(root_path, subdir_name)
    # 判断是否是文件夹
    if os.path.isdir(subdir_path):
        file_list = os.listdir(subdir_path)  # subdir_path文件夹下的所有jpg和json文件
        for file_name in file_list:
            file_path = os.path.join(subdir_path, file_name)
            pic_name = ''
            ext = ''
            file_name_split = file_name.split('.')
            if len(file_name_split) > 2:
                # such as louisvuitton_2905806368.jpg_400x480.jpg
                ext = file_name_split[-1]
                for i in range(len(file_name_split) - 1):
                    pic_name += '.'
                    pic_name += file_name_split[i]
                pic_name = pic_name[1:]
            else:
                pic_name, ext = file_name_split

            # 图片
            if ext in image_ext:
                shutil.copy(file_path, new_root_path)
            # json文件
            elif ext == json_ext:
                with open(file_path, 'r', encoding='utf-8') as f:
                    pic_json_dict = json.load(f)
                imageHeight = pic_json_dict["imageHeight"]
                imageWidth = pic_json_dict["imageWidth"]

                # 后缀名匹配
                for file_name_iter in file_list:
                    pic_name_iter = ''
                    ext_iter = ''
                    file_name_iter_split = file_name_iter.split('.')
                    if len(file_name_iter_split) > 2:
                        ext_iter = file_name_iter_split[-1]
                        for i in range(len(file_name_iter_split) - 1):
                            pic_name_iter += '.'
                            pic_name_iter += file_name_iter_split[i]
                        pic_name_iter = pic_name_iter[1:]
                    else:
                        pic_name_iter, ext_iter = file_name_iter_split

                    if pic_name_iter == pic_name and ext_iter in image_ext:
                        pic_file_name = pic_name + "." + ext_iter

                # 向train_dict的“images”模块写入信息
                train_dict_images_dict = {}
                train_dict_images_dict["height"] = imageHeight
                train_dict_images_dict["width"] = imageWidth
                train_dict_images_dict["id"] = train_dict_images_id
                train_dict_images_dict["file_name"] = pic_file_name
                train_dict["images"].append(train_dict_images_dict)

                # 向train_dict的“annotations”模块写入信息
                bboxs = pic_json_dict["shapes"]
                for bbox in bboxs:
                    train_dict_annotations_dict = {}
                    train_dict_annotations_dict["id"] = train_dict_annotations_id
                    train_dict_annotations_id += 1
                    train_dict_annotations_dict["image_id"] = train_dict_images_id

                    for category in categories_list:
                        if subdir_name == category['name']:
                            train_dict_annotations_dict["category_id"] = category['id']

                    # 注意x和y的顺序
                    bbox_x0 = int(bbox["points"][0][0])
                    bbox_y0 = int(bbox["points"][0][1])
                    bbox_x1 = int(bbox["points"][1][0])
                    bbox_y1 = int(bbox["points"][1][1])
                    bbox_w = bbox_x1 - bbox_x0
                    bbox_h = bbox_y1 - bbox_y0
                    train_dict_annotations_dict["bbox"] = [bbox_x0, bbox_y0, bbox_w, bbox_h]

                    train_dict_annotations_dict["iscrowd"] = 0
                    train_dict_annotations_dict["area"] = bbox_w * bbox_h

                    train_dict["annotations"].append(train_dict_annotations_dict)

                train_dict_images_id += 1
    # print(f"subdir {subdir_name} 's files write to json file succeeded")

# 将train_dict写入json文件
train_annoations_path = "logo_det0526.json"
with open(train_annoations_path, 'w', encoding='utf-8') as f:
    json.dump(train_dict, f, indent=4, ensure_ascii=False)


# 在训练集图片上标注box
def visual(img_path, bbox, cls_id, class_names):
    img = cv2.imread(img_path)
    color = (_COLORS[cls_id] * 255).astype(np.uint8).tolist()
    text = '{}'.format(class_names[cls_id - 1])
    txt_color = (0, 0, 0) if np.mean(_COLORS[cls_id]) > 0.5 else (255, 255, 255)
    font = cv2.FONT_HERSHEY_SIMPLEX
    txt_size = cv2.getTextSize(text, font, 0.4, 1)[0]
    txt_bk_color = (_COLORS[cls_id] * 255 * 0.7).astype(np.uint8).tolist()

    x_min = bbox[0]
    y_min = bbox[1]
    width = bbox[2]
    height = bbox[3]
    cv2.rectangle(img, (x_min, y_min), (x_min + width, y_min + height), color, thickness=2)
    cv2.rectangle(img, (x_min, y_min - int(1.5 * txt_size[1])), (x_min + txt_size[0] + 1, y_min), txt_bk_color, -1)
    cv2.putText(img, text, (x_min, y_min - 3), font, 0.4, txt_color, thickness=1)

    cv2.imwrite(img_path, img)


# 标注gt框
if args.visual:
    # 新建文件夹
    visual_dir_path = 'visual_dir'
    isExists = os.path.exists(visual_dir_path)
    if isExists:
        shutil.rmtree(visual_dir_path)
    os.mkdir(visual_dir_path)

    # 复制图片
    file_list = os.listdir(new_root_path)
    for file_name in file_list:
        old_file_path = os.path.join(new_root_path, file_name)
        shutil.copy(old_file_path, visual_dir_path)

    # 可视化
    files = os.listdir(visual_dir_path)
    images_list = train_dict["images"]
    annotations_list = train_dict["annotations"]
    CLASSES = categories_list

    for image_name in tqdm(files, desc="visual process"):
        image_name = os.path.join(visual_dir_path, image_name)
        for i in range(len(images_list)):
            pic_file_name = os.path.basename(image_name)
            if pic_file_name == images_list[i]["file_name"]:
                image_id = images_list[i]["id"]
                for j in range(len(annotations_list)):
                    if image_id == annotations_list[j]["image_id"]:
                        bbox = annotations_list[j]["bbox"]
                        cls_id = annotations_list[j]["category_id"]
                        visual(image_name, bbox, cls_id, CLASSES)