yolo格式数据标签转coco格式

最新推荐文章于 2024-07-05 09:30:53 发布

carambola_TaoYang

最新推荐文章于 2024-07-05 09:30:53 发布

阅读量6.4k

点赞数 5

文章标签： 1024程序员节

本文链接：https://blog.csdn.net/carambola_/article/details/127499615

版权

本文介绍了如何将Yolo格式的物体检测标签转换为Coco格式，以便于使用Deformable DETR等需要Coco格式数据的模型。转换过程包括两个步骤：首先，将Yolo的bbox标签转换为Coco格式，并保存为txt文件；其次，生成json注释文件，包括图像信息和类别信息。转换过程中涉及到图像尺寸调整、类别映射以及 bbox 数据结构的转换。

摘要由CSDN通过智能技术生成

yolo格式数据标签转coco格式

最近在使用Deformable detr时输入数据格式为coco类型的标签，因此根据别人的代码改写了一下，将yolo类型标签转换为coco格式。
总共分为两步：
1. 第一步：将yolo的box标签类型（xMin, yMim, xMax, yMax）转化为coco标签类型（xMin, yMim, xMax, yMax），并保存为txt文件，同时相较于yolo的txt文件，coco的注释文件中还多了一列标签对应的图像名称。
输入参数：
（1） 原始标签路径originLabelsDir：
存放标签的文件夹：

存放标签文件夹

txt文件中的内容
因为我这边只有一类所以标签就全是0

在这里插入图片描述
（2）原始图片路径originImagesDir：

在这里插入图片描述
(3) 生成的临时注释txt文件
生成的临时文件中的内容应该如下：
将yolo的标签全部存入一个txt文件中
将yolo格式的标签：classId, xCenter, yCenter, w, h转换为
coco格式：classId, xMin, yMim, xMax, yMax格式
coco的id编号从1开始计算，所以这里classId应该从1开始计算
最终annos.txt中每行为imageName, classId, xMin, yMim, xMax, yMax, 一个bbox对应一行
在这里插入图片描述
代码如下

import os
import cv2
import json
from tqdm import tqdm

def yolo2txt(originLabelsDir, originImagesDir, saveTempTxt, img_suffix='png'):
    """
    将yolo的标签全部存入一个txt文件中
    将yolo格式的标签：classId, xCenter, yCenter, w, h转换为
    coco格式：classId, xMin, yMim, xMax, yMax格式
    coco的id编号从1开始计算，所以这里classId应该从1开始计算
    最终annos.txt中每行为imageName, classId, xMin, yMim, xMax, yMax, 一个bbox对应一行
    originLabelsDir: 原始yolo标签路径
    originImagesDir: 原始图像路径
    ssaveTempTxt: 保存txt路径
    )
    """
    txtFileList = os.listdir(originLabelsDir)
    print(f"image number is {len(txtFileList)}")
    with open(saveTempTxt, 'w') as fw:
        for txtFile in tqdm(txtFileList, desc="generating COCO format"):
            # 读取图像长宽
            imagePath = os.path.join(originImagesDir,
                                     txtFile.replace('txt', img_suffix))
            assert os.path.exists(imagePath), f"can\'t find this image {imagePath}"
            image = cv2.imread(imagePath)
            H, W, _ = image.shape

            with open(os.path.join(originLabelsDir, txtFile), 'r') as fr:
                labelList = fr.readlines()
                for label in labelList:
                    label = label.strip().split()
                    x = float(label[1])
                    y = float(label[2])
                    w = float(label[3])
                    h = float(label[4])

                    # convert x,y,w,h to x1,y1,x2,y2
                    x1 = (x - w / 2) * W
                    y1 = (y - h / 2) * H
                    x2 = (x + w / 2) * W
                    y2 = (y + h / 2) * H
                    # 为了与coco标签方式对，标签序号从1开始计算
                    fw.write(txtFile.replace('txt', img_suffix) + ' {} {} {} {} {}\n'.format(int(label[0]) + 1, x1, y1, x2, y2))
if __name__ == "__main__":
    # 原始标签路径
    originLabelsDir = '/home/taoyang/Documents/datasets/datasets/wheat_panicle_dataset/train/labels'
    # 转换后的文件保存路径
    saveTempTxt = '/home/taoyang/Documents/datasets/datasets/wheat_panicle_dataset/train/annos_temp.txt'
    # 原始标签对应的图片路径
    originImagesDir = '/home/taoyang/Documents/datasets/datasets/wheat_panicle_dataset/images/train2017'
    yolo2txt(originLabelsDir, originImagesDir, saveTempTxt)

第二步：生成json注释文件，并保存。
有三个输入分别为：
原始图像保存的文件夹：originImagesDir
类别txt文件：classtxt

上一步生成的临时注释文件：annostxt

最终生成的注释文件.json保存在与annotxt相同路径下

第二步代码如下：

def yolo2json(originImagesDir, classtxt, annostxt):
    # dataset用于保存所有数据的图片信息和标注信息
    dataset = {'categories': [], 'annotations': [], 'images': []}

    # 打开类别标签
    with open(classtxt, 'r') as f:
        classes_ori = f.readlines()
        classes = [i.strip().split(' ') for i in classes_ori if i.strip() != '']

    # 建立类别标签和数字id的对应关系
    for i, cls in classes:
        dataset['categories'].append({'id': i, 'name': cls})

    # 读取images文件夹的图片名称
    indexes = os.listdir(originImagesDir)

    anno_id = -1
    # 读取Bbox信息
    with open(annostxt) as tr:
        annos = tr.readlines()
        # ---------------接着将，以上数据转换为COCO所需要的格式---------------
        with tqdm(total=len(indexes)) as pbar:
            for k, index in enumerate(indexes):
                # 用opencv读取图片，得到图像的宽和高
                im = cv2.imread(os.path.join(originImagesDir, index))
                assert im.all() != None, f"can\'t find this image {os.path.join(originImagesDir, index)}"
                height, width, _ = im.shape
                # 添加图像的信息到dataset中
                dataset['images'].append({'file_name': index,
                                          'id': k,
                                          'width': width,
                                          'height': height})
                del_annos = [] # 保存已经匹配到的注释信息
                for anno in annos:
                    parts = anno.strip().split()

                    # 如果图像的名称和标记的名称对上，则添加标记
                    if parts[0] == index:
                        del_annos.append(anno)
                        anno_id += 1
                        # 类别
                        cls_id = parts[1]
                        # x_min
                        x1 = float(parts[2])
                        # y_min
                        y1 = float(parts[3])
                        # x_max
                        x2 = float(parts[4])
                        # y_max
                        y2 = float(parts[5])
                        width = x2 - x1
                        height = y2 - y1
                        assert width > 0 and height > 0, f"width or height of {index}\'s box is not positive"
                        dataset['annotations'].append({
                            'area': width * height,
                            'bbox': [x1, y1, width, height],
                            'category_id': cls_id, # 根据自己的后续代码决定是否转化为int类型，原来的代码这边是转化成int类
                            # 型，但是我后面的代码中category_id用的是str类型，导致我一直读不进对应标签的图像
                            'id': anno_id,
                            'image_id': k,
                            'iscrowd': 0,
                            # 左上角以及右下角坐标
                            'segmentation': [[x1, y1, x2, y2]]
                        })

                # 删除匹配到的注释
                for da in del_annos:
                    annos.remove(da)
                pbar.update(1)

    if len(annos) != 0:
        print(f"\033[31m can\'t match image for these annotations:\n{annos}\033[0m")

    # 保存结果的文件夹， 保存在与annostxt相同文件夹下，命名为annotations.json
    savefile = os.path.join(annostxt.rsplit('/', 1)[-2], 'annotations.json')
    with open(savefile, 'w', encoding='utf-8') as f:
        json.dump(dataset, f, ensure_ascii=False, indent=1)
if __name__ == "__main__":
    originImagesDir = '/home/taoyang/Documents/datasets/datasets/wheat_panicle_dataset/images/train2017'
    classtxt = '/home/taoyang/Documents/datasets/datasets/wheat_panicle_dataset/train/classes.txt'
    annostxt = '/home/taoyang/Documents/datasets/datasets/wheat_panicle_dataset/train/annos_temp.txt'
    yolo2json(originImagesDir, classtxt, annostxt)