语义分割json转目标检测txt标签

下面代码写于labelme标记不规则图形（规则也行）产生的json文件，转yolo（这里是yolov5）的txt标签文件，避免再次标记

import json
import os
from pathlib import Path
import re
from tqdm import tqdm
import shutil


def convert_annotation(image_id, paths):
    global label_map

    def find_box(points):  # 该函数用来找出xmin, xmax, ymin ,ymax 即bbox包围框
        _x, _y = [float(pot[0]) for pot in points], [float(pot[1]) for pot in points]
        return min(_x), max(_x), min(_y), max(_y)

    def convert(size, bbox):  # 转为中心坐标
        # size: (原图宽， 原图长)
        center_x, center_y = (bbox[0] + bbox[1]) / 2.0 - 1, (bbox[2] + bbox[3]) / 2.0 - 1
        center_w, center_h = bbox[1] - bbox[0], bbox[3] - bbox[2]
        return center_x / size[0], center_y / size[1], center_w / size[0], center_h / size[1]

    final_label_path, final_output_path = paths
    label_json_url = os.path.join(final_label_path, f'{image_id}.json')
    # 输出到 ：final_output_path / f'{image_id}_leftImg8bit.txt'

    load_dict = json.load(open(label_json_url, 'r'))  # 图像的实例
    output_cache = []
    for obj in load_dict['shapes']:  # load_dict['objects'] -> 目标的几何框体
        obj_label = obj['label']  # 目标的类型
        if obj_label not in list(label_map.keys()):
            continue

        if obj_label not in label_map.keys():  # 记录目标类型转为int值
            label_map[obj_label] = len(label_map.keys())  # 标签从0开始

        x, y, w, h = convert((load_dict['imageWidth'], load_dict['imageHeight']), find_box(obj['points']))  # 归一化为中心点

        # yolo 标准格式：img.jpg -> img.txt
        # 内容的类别 归一化后的中心点x坐标 归一化后的中心点y坐标 归一化后的目标框宽度w 归一化后的目标况高度h
        output_cache.append(f'{label_map[obj_label]} {x} {y} {w} {h}\n')

    # with open(final_output_path / f'{image_id}_leftImg8bit.txt', 'w') as label_f:  # 写出标签文件
    with open(os.path.join(final_output_path, image_id + ".txt"), 'w') as label_f:  # 写出标签文件
        label_f.writelines(output_cache)


def mkdir(url):
    if not os.path.exists(url):
        os.makedirs(url)


if __name__ == '__main__':
    root_dir = r'D:\BaiduNetdiskDownload\laser_spot_yolo_dataset'
    image_dir = os.path.join(root_dir, 'laser_img')
    label_dir = os.path.join(root_dir, 'laser_json')
    image_output_root_dir = os.path.join(root_dir, 'images')
    label_output_root_dir = os.path.join(root_dir, 'labels')
    image_label = ".jpg"
    label_map = { "1": 0
   }  # 存放所有的类别标签  eg. {'car': 0, 'person': 1}
    for _t_ in tqdm(os.listdir(image_dir)):  # _t_ as ['train', 'test' 'val;]
        type_files = []  # 存放一类的所有文件，如训练集所有文件
        mkdir(os.path.join(image_output_root_dir , _t_)), mkdir(os.path.join(label_output_root_dir, _t_))
        # for cities_name in os.listdir(os.path.join(image_dir, _t_)):
        #     _final_img_path = os.path.join(image_dir, _t_, cities_name)  # root_dir / laser_img / test / berlin
        #     _final_label_path = os.path.join(label_dir, _t_, cities_name)  # root_dir / laser_json / test / berlin
        _final_img_path = os.path.join(image_dir, _t_)  # root_dir / laser_img / test / berlin
        _final_label_path = os.path.join(label_dir, _t_)  # root_dir / laser_json / test / berlin

        # berlin_000000_000019_leftImg8bit.png -> berlin_000000_000019_gtFine_polygons.json
        # map()会根据提供的函数对指定序列做映射。
        # 第一个参数function以参数序列中的每一个元素调用function函数，返回包含每次function函数返回值的新列表。
        image_ids = list(map(lambda s: re.sub(f'{image_label}', '', s), os.listdir(_final_img_path)))
        # image_ids = os.listdir(_final_img_path)
        # print(names[:0])  -> berlin_000000_000019

        for img_id in image_ids:
            convert_annotation(img_id, [_final_label_path, os.path.join(label_output_root_dir, _t_)])  # 转化标签
            # 接下来几个步骤是为准备yolov5数据集格式
            # img_file = f'{img_id}_leftImg8bit.png'
            img_file = img_id + image_label
            shutil.copy(os.path.join(_final_img_path, img_file), f'{root_dir}/images/{_t_}/{img_file}')  # 复制移动图片
            type_files.append(f'images/{_t_}/{img_id}{image_label}\n')

        with open(os.path.join(root_dir, f'yolo_{_t_}.txt'), 'w') as f:  # 记录训练样本等的具体内容
            f.writelines(type_files)

    # with open(label_output_root_dir / 'classes.txt', 'w') as f:  # 写出类别对应
    #     for k, v in label_map.items():
    #         f.write(f'{k}\n')
    # print([k for k in label_map.keys()], len([k for k in label_map.keys()]))