CV图像处理小工具——语义分割json生成检测框json

语义分割json生成检测框json

import json
import os
from os import listdir, getcwd
from os.path import join
import os.path


rootdir = 'F:/dataset/'# 写自己存放图片的数据地址
input_dir = 'F:/dataset/labels_json/'
output_dir = 'F:/dataset/labels_box/'
def position(pos):
    # 该函数用来找出xmin,ymin,xmax,ymax即bbox包围框
    x = []
    y = []
    nums = len(pos)
    for i in range(nums):
        x.append(pos[i][0])
        y.append(pos[i][1])
    x_max = max(x)
    x_min = min(x)
    y_max = max(y)
    y_min = min(y)
    b = (float(x_min), float(x_max), float(y_min), float(y_max))
    return b




def convert_annotation(image_ids, input_dir, output_dir):
    for image_id in image_ids:
        print(image_id)
        input_file_path = os.path.join(input_dir, f"{image_id}.json")
        output_file_path = os.path.join(output_dir, f"{image_id}.json")

        try:
            with open(input_file_path, 'r') as load_f:
                load_dict = json.load(load_f)

            w = load_dict['imageWidth']
            h = load_dict['imageHeight']
            objects = load_dict['shapes']

            annotations = []  # 创建一个空列表来存储注解

            for obj in objects:
                # 去除可能的额外空格，并设置默认标签（如果需要）
                labels = obj['label'].strip()
                # 根据标签处理形状
                if labels in [
                    "class1", "class2",
                    "class3", "class4"
                ]:
                    pos = obj['points']
                    b = position(pos)
                    cls_id_mapping = {
                        "class1": 1,
                        "class2": 2,
                        "class3": 3,
                        "class4": 4,
                        
                    }
                    cls_id = cls_id_mapping[labels]
                    annotation = {
                        'label': labels,
                        'class_id': cls_id,
                        'bbox': b
                    }
                    annotations.append(annotation)

            output_dict = {
                "version": "5.0.2",
                "flags": {},
                "shapes": [],  # shapes数组为空，因为注解信息已经放在annotations中
                "imagePath": load_dict['imagePath'],
                "imageData": None,
                "imageHeight": h,
                "imageWidth": w,
                "annotations": annotations  # 添加注解列表
            }
            for shape in load_dict['shapes']:
                # 假设每个多边形都有四个点，我们可以直接取对角线上的两个点来定义矩形
                # 这里我们取第一个点和第三个点（或者您可以根据具体情况选择其他点对）
                objects = load_dict['shapes']
                for obj in objects:
                    # 去除可能的额外空格，并设置默认标签（如果需要）
                    labels = obj['label'].strip()
                    if labels in [
                       "class1", "class2",
                    "class3", "class4"
                    ]:
                        pos = obj['points']
                        b = position(pos)
                        rect_points = [
                            [b[0], b[3]],
                            [b[1], b[2]]
                            ]

                        new_shape = {
                        "label": labels,  # 保留原多边形的标签（或者您可以根据需要生成新的标签）
                        "points": rect_points,  # 使用调整后的点来表示矩形（但这里我们实际上只使用了两个点）
                        # 注意：由于我们简化了问题，并没有真正地使用四个点来定义一个完整的矩形
                        # 在实际应用中，您可能需要更精确地处理这个问题。
                        "group_id": None,  # 您可以根据需要设置这个字段的值
                        "shape_type": "rectangle",  # 指定形状类型为矩形
                        "flags": {}  # 保留空的flags字段（或者您可以根据需要填充它）
                    }
                        output_dict['shapes'].append(new_shape)

                # 写入新的JSON文件
            with open(output_file_path, 'w') as output_f:
                json.dump(output_dict, output_f, indent=4)

        except Exception as e:
            print(f"Error processing {input_file_path}: {e}")


def image_id(rootdir):
    a = []
    for parent, dirnames, filenames in os.walk(rootdir):
        for filename in filenames:
            # print(filename)
            if filename.endswith('.jpg'):
                filename = os.path.splitext(filename)[0]
            if filename.endswith('.jpeg'):
                filename = os.path.splitext(filename)[0]
            if filename.endswith('.JPG'):
                filename = os.path.splitext(filename)[0]
            if filename.endswith('.JPEG'):
                filename = os.path.splitext(filename)[0]
            a.append(filename)
    return a


names = image_id("F:/dataset/images/"),

for image_id in names:
    convert_annotation(image_id, input_dir, output_dir)