代码收藏家技术教程 2024-10-26

Python制作Coco语义分割数据集标签，已解决标签类别重复标注问题【超简单！】

我自己制作好了一套COCO2017语义分割数据集，一共122218张图片。（试过上传到到百度网盘上面，结果文件太大了，龟速0_0，就放弃了）大家可以根据我下面提供的代码制作coco标签。

一、官网下载文件

COCO官方没有提供标签图像，在官网上只能下载到原图和json文件，后续还需要用到python代码来制作标签图像。（coco官方标注的标签图像是会有类别重复标注这种情况的，这可能就是他为什么没有上传标签图像的原因之一吧）

进入官网COCO – Common Objects in Context (cocodataset.org)然后下载这3个文件：

二、python制作标签

根据博文从coco2017数据集中提取语义分割mask-CSDN博客的方法，处理后的coco标签会出现标签类别重复标注的问题，也即是说标签中的类别（加上背景）不是81类，实际比81类要多，这样标签就没法用了。

他的方法存在的问题：

1、语义分割标签是一个像素值对应一个类别，正常情况下标签像素值应该是从0-81，像素值超过了81的像素点在模型训练时候读取不了，这些像素点相当于没有进行标注，这些像素区域的特征也就没有学习到。

2、按照他的方法处理，前81个像素值对应的类别也会出错（标签类别重复标注的点会出现像id为2和id为65的进行相加，最终id就变成了67，但是67是其他类别，相应地学习到的特征就会出错），模型学习出来的东西也不知道在学什么，相当混乱，没法用来训练自己的模型。

3、对于类别重复标注的像素区域，如果采用小号id覆盖大id这种方法的话，会出现桌子把橙子、甜甜圈覆盖掉的问题，被覆盖的物体这一类别就直接没了！

根据他的python代码，我重新定义了物体类别所对应的id，修改了选择类别的方式，对类别重复标注的像素区域，用大id类别覆盖小id类别，不用出现像蛋糕被桌子覆盖这种问题了，摆在上面的物体不会被下面的物体覆盖，前面位置的物体不会被后面的物体覆盖，这样就符合现实中物体摆放的逻辑了。

文件结构如下图设置：

在COCOdataset文件夹下有 annotations 和 images 两个文件夹。annotations中放的是json文件， images中放的是train2017和val2017两个文件夹

文件结构如下图设置：

在COCOdataset文件夹下有 annotations 和 images 两个文件夹。annotations中放的是json文件， images中放的是train2017和val2017两个文件夹

D:/COCO2017DataSet
          |  
          |--COCOdataset
                  |
                  |--annotations
                        |
                        |--instances_train2017.json
                        |--instances_val2017.json
                  |--images
                        |
                        |--train2017（这里放的是训练集jpg原图）
                        |--val2017（这里放的是验证集jpg原图）

from PIL import Image
import imgviz
import argparse
import os
import tqdm
from pycocotools.coco import COCO
import shutil
import numpy as np

"""
自定义图片中像素点的类别id，制作mask时可以对这些id进行转换
原图类别id  转换为  自己定义的类别id
注意原图的类别id并不连续 但总数为80

此代码对于类别重复标注的像素点，设置为id最大那一类
"""
id_mapping = {
    1:2,
    2:10,
    3:9,
    4:8,
    5:7,
    6:6,
    7:5,
    8:4,
    9:3,
    10:20,
    11:19,
    13:18,
    14:17,
    15:27,
    16:30,
    17:32,
    18:31,
    19:1,
    20:16,
    21:15,
    22:14,
    23:13,
    24:12,
    25:11,
    27:24,
    28:47,
    31:35,
    32:33,
    33:34,
    34:48,
    35:46,
    36:45,
    37:52,
    38:44,
    39:51,
    40:50,
    41:43,
    42:42,
    43:49,
    44:75,
    46:74,
    47:76,
    48:77,
    49:79,
    50:78,
    51:54,
    52:73,
    53:80,
    54:57,
    55:73,
    56:72,
    57:71,
    58:70,
    59:55,
    60:69,
    61:56,
    62:28,
    63:26,
    64:53,
    65:21,
    67:25,
    70:22,
    72:36,
    73:41,
    74:62,
    75:63,
    76:61,
    77:64,
    78:40,
    79:39,
    80:38,
    81:37,
    82:23,
    84:29,
    85:58,
    86:65,
    87:67,
    88:59,
    89:60,
    90:66
}


def save_colored_mask(mask, save_path):
    lbl_pil = Image.fromarray(mask.astype(np.uint8), mode="P")
    colormap = imgviz.label_colormap()
    lbl_pil.putpalette(colormap.flatten())
    lbl_pil.save(save_path)


def main(args):
    annotation_file = os.path.join(args.input_dir, 'annotations', 'instances_{}.json'.format(args.split))
    os.makedirs(os.path.join(args.input_dir, 'SegmentationClass'), exist_ok=True)
    os.makedirs(os.path.join(args.input_dir, 'JPEGImages'), exist_ok=True)

    coco = COCO(annotation_file)    # 加载COCO标签的注释文件

    catIds = coco.getCatIds()
    # 获取 标 签 中所有类别的ID 返回：catIds是一个列表，包含了所有类别的ID。每个ID对应一个类别，例如 {'person': 1, 'car': 3, ...}。
    # 细节：如果不指定任何参数（如catIds），coco.getCatIds()会返回所有类别的ID。

    imgIds = coco.getImgIds()
    # 作用：获取 标签 中所有图像的ID。
    # 返回：imgIds是一个列表，包含了所有图像的ID。每个ID对应一个图像在注释文件中的唯一标识符。
    # 细节：如果不指定任何参数（如imgIds），coco.getImgIds()会返回数据集中所有图像的ID。

    print("catIds len:{}, imgIds len:{}".format(len(catIds), len(imgIds)))

    for imgId in tqdm.tqdm(imgIds, ncols=100):
        # 遍历图像id
        # tqdm.tqdm(imgIds, ncols=100)：使用 tqdm库显示处理进度条，imgIds是所有图像的ID列表，ncols=100设置进度条的宽度为100列。

        img = coco.loadImgs(imgId)[0]
        # 根据图像ID加载图像信息

        annIds = coco.getAnnIds(imgIds=img['id'], catIds=catIds, iscrowd=None)
        # 获取图像的注释ID

        anns = coco.loadAnns(annIds)
        # coco.loadAnns(annIds)：根据注释ID加载注释信息。返回一个包含所有注释信息的列表。

        if len(annIds) > 0:

            # 初始化掩膜和类别字典： 先初始化掩膜为全0，并使用一个字典来存储每个像素的最大类别ID
            # 更新掩膜和类别字典： 遍历每个标注，更新字典中的最大类别ID，并在掩膜中保存这个最大类别ID
            # 将类别字典转换为掩膜： 将最终的类别字典转换为掩膜，并保存结果

            mask = np.zeros((img['height'], img['width']), dtype=np.uint8)
            class_id_map = np.zeros((img['height'], img['width']), dtype=np.uint8)

            for ann in anns:
                category_id = ann['category_id']
                me_id = id_mapping[category_id]
                ann_mask = coco.annToMask(ann)

                # 更新类别字典和掩膜
                class_id_map[ann_mask > 0] = np.maximum(class_id_map[ann_mask > 0], me_id)
                mask = class_id_map

            img_origin_path = os.path.join(args.input_dir, 'images', args.split, img['file_name'])
            img_output_path = os.path.join(args.input_dir, 'JPEGImages', img['file_name'])
            seg_output_path = os.path.join(args.input_dir, 'SegmentationClass',
                                           img['file_name'].replace('.jpg', '.png'))
            shutil.copy(img_origin_path, img_output_path)
            save_colored_mask(mask, seg_output_path)


def get_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("--input_dir", default="D:/COCO2017DataSet/COCOdataset", type=str,
                        help="input dataset directory")
    parser.add_argument("--split", default="train2017", type=str,
                        help="train2017 or val2017")
    return parser.parse_args()


if __name__ == '__main__':
    args = get_args()
    main(args)

在修改时把default修改为你自己的路径：default=“你自己的路径”，

修改id_mapping，可以给对原始的物体类别进行自定义id设置。

在envi中随便打开一张我们处理好的标签，可以看到人这一类别的像素值已经正确赋值为了“2”，而且没有出现类被重复标注的情况。其他物体的类别对照下表查看。

序号	原始id	原始id对应的物体	自定义id	按自定义id排序	自定义id对应的物体	像素值	类别
1	1	人	2	1	马	1	马
2	2	自行车	10	2	人	2	人
3	3	汽车	9	3	船	3	船
4	4	摩托车	8	4	卡车	4	卡车
5	5	飞机	7	5	火车	5	火车
6	6	公交车	6	6	公交车	6	公交车
7	7	火车	5	7	飞机	7	飞机
8	8	卡车	4	8	摩托车	8	摩托车
9	9	船	3	9	汽车	9	汽车
10	10	红绿灯	20	10	自行车	10	自行车
11	11	消防栓	19	11	长颈鹿	11	长颈鹿
12	13	停止标志	18	12	斑马	12	斑马
13	14	停车收费表	17	13	熊	13	熊
14	15	长凳	27	14	象	14	象
15	16	鸟	30	15	牛	15	牛
16	17	猫	32	16	羊	16	羊
17	18	狗	31	17	停车收费表	17	停车收费表
18	19	马	1	18	停止标志	18	停止标志
19	20	羊	16	19	消防栓	19	消防栓
20	21	牛	15	20	红绿灯	20	红绿灯
21	22	象	14	21	床	21	床
22	23	熊	13	22	马桶	22	马桶
23	24	斑马	12	23	冰箱	23	冰箱
24	25	长颈鹿	11	24	背包	24	背包
25	27	背包	24	25	餐桌	25	餐桌
26	28	雨伞	47	26	长椅	26	长椅
27	31	手提包	35	27	长凳	27	长凳
28	32	领带	33	28	椅子	28	椅子
29	33	手提箱	34	29	书	29	书
30	34	飞盘	48	30	鸟	30	鸟
31	35	滑雪板	46	31	狗	31	狗
32	36	单板滑雪	45	32	猫	32	猫
33	37	运动球	52	33	领带	33	领带
34	38	风筝	44	34	手提箱	34	手提箱
35	39	棒球棒	51	35	手提包	35	手提包
36	40	棒球手套	50	36	电视	36	电视
37	41	滑板	43	37	洗碗槽	37	洗碗槽
38	42	冲浪板	42	38	烤面包机	38	烤面包机
39	43	网球拍	49	39	烤箱	39	烤箱
40	44	瓶子	75	40	微波炉	40	微波炉
41	46	红酒杯	74	41	笔记本电脑	41	笔记本电脑
42	47	杯子	76	42	冲浪板	42	冲浪板
43	48	叉子	77	43	滑板	43	滑板
44	49	刀	79	44	风筝	44	风筝
45	50	勺	78	45	单板滑雪	45	单板滑雪
46	51	碗	54	46	滑雪板	46	滑雪板
47	52	香蕉	73	47	雨伞	47	雨伞
48	53	苹果	80	48	飞盘	48	飞盘
49	54	三明治	57	49	网球拍	49	网球拍
50	55	橙子	73	50	棒球手套	50	棒球手套
51	56	西兰花	72	51	棒球棒	51	棒球棒
52	57	胡萝卜	71	52	运动球	52	运动球
53	58	热狗	70	53	盆栽	53	盆栽
54	59	比萨	55	54	碗	54	碗
55	60	甜甜圈	69	55	比萨	55	比萨
56	61	蛋糕	56	56	蛋糕	56	蛋糕
57	62	椅子	28	57	三明治	57	三明治
58	63	长椅	26	58	时钟	58	时钟
59	64	盆栽	53	59	泰迪熊	59	泰迪熊
60	65	床	21	60	吹风机	60	吹风机
61	67	餐桌	25	61	键盘	61	键盘
62	70	马桶	22	62	鼠标	62	鼠标
63	72	电视	36	63	遥控器	63	遥控器
64	73	笔记本电脑	41	64	手机	64	手机
65	74	鼠标	62	65	花瓶	65	花瓶
66	75	遥控器	63	66	牙刷	66	牙刷
67	76	键盘	61	67	剪刀	67	剪刀
68	77	手机	64	69	甜甜圈	69	甜甜圈
69	78	微波炉	40	70	热狗	70	热狗
70	79	烤箱	39	71	胡萝卜	71	胡萝卜
71	80	烤面包机	38	72	西兰花	72	西兰花
72	81	洗碗槽	37	73	香蕉	73	香蕉
73	82	冰箱	23	73	橙子	73	橙子
74	84	书	29	74	红酒杯	74	红酒杯
75	85	时钟	58	75	瓶子	75	瓶子
76	86	花瓶	65	76	杯子	76	杯子
77	87	剪刀	67	77	叉子	77	叉子
78	88	泰迪熊	59	78	勺	78	勺
79	89	吹风机	60	79	刀	79	刀
80	90	牙刷	66	80	苹果	80	苹果
						0	黑色背景