【python脚本】coco数据集格式转labelme格式以及筛选指定类别的数据

目标任务：通过coco官方数据集转成实例分割用labelme格式，并且筛选出指定的类别的image和json文件

一、数据集准备

1. 官方数据集 https://cocodataset.org/#download

官方下载比较慢，不推荐

2. 第三方下载 https://blog.csdn.net/ji_meng/article/details/124959983?spm=1001.2014.3001.5506

如果对数据集数量没有太大要求，推荐用val2017，比较小，下载速度快。

axel -n 100 http://images.cocodataset.org/zips/val2017.zip

axel -n 100 http://images.cocodataset.org/annotations/annotations_trainval2017.zip

├── annotations_trainval2017
├── annotations_trainval2017.zip
├── val2017
└── val2017.zip

二、数据集预处理

本案例是通过实例分割的情况，需要对原始数据做一下预处理，需要将coco格式转换成labelme格式

1. coco2labelme.py

# -*- coding: utf-8 -*-
import glob
import os
import cv2
import json
import io

"""
要将官方的coco数据集转成labelme，稍微有点不同，因为官方的coco的实际标签是排列到90，当将coco数据集转成labelme数据集时候，需要修改一下标签，插入空标签

分别要在（从１开始）　12 26 29 30 45 66 68 69 71 83 插入空labelme
"""

coco=["person", "bicycle", "car", "motorcycle", "airplane", "bus", "train", "truck", "boat", "traffic light",
        "fire hydrant","" ,"stop sign", "parking meter", "bench", "bird", "cat", "dog", "horse", "sheep", "cow",
        "elephant", "bear", "zebra", "giraffe", "","backpack", "umbrella","", "","handbag", "tie", "suitcase", "frisbee",
        "skis", "snowboard", "sports ball", "kite", "baseball bat", "baseball glove", "skateboard", "surfboard",
        "tennis racket", "bottle","", "wine glass", "cup", "fork", "knife", "spoon", "bowl", "banana", "apple",
        "sandwich", "orange", "broccoli", "carrot", "hot dog", "pizza", "donut", "cake", "chair", "couch",
        "potted plant", "bed","", "dining table", "","","toilet", "","tv", "laptop", "mouse", "remote", "keyboard", "cell phone",
        "microwave", "oven", "toaster", "sink", "refrigerator", "book", "","clock", "vase", "scissors", "teddy bear",
        "hair drier", "toothbrush"]
 
label=dict()
for idx,item in enumerate(coco):
    label.update({idx:item})

labelme_json_img_path= "/home/ubuntu/pick_datasets/coco_val_all"    # 保存数据集
coco_json_path = '/home/ubuntu/pick_datasets/annotations_trainval2017/annotations'  # coco annotations
jpg_path= "/home/ubuntu/pick_datasets/val2017"  # coco val images
coco_json=glob.glob(os.path.join(coco_json_path,"*.json"))[2]
# print(coco_json)
file_json = io.open(coco_json,'r',encoding='utf-8')
m_json_data = file_json.read()
m_data = json.loads(m_json_data)
#m_type=m_data['type']
 
for item in m_data['images']:
    m_images_file_name = item['file_name']
    (filename_path, m_filename) = os.path.split(m_images_file_name)
    (m_name, extension) = os.path.splitext(m_filename)
    m_image=cv2.imread(os.path.join(jpg_path,m_name+".jpg"))
    m_images_height = item['height']
    m_images_width = item['width']
    m_images_id = item['id']
    data = {}
    data['version'] = "5.1.1"
    data['flags'] = {}
    data["shapes"] = []
    for annit in m_data['annotations']:
        m_image_id=annit['image_id']
        m_category_id=annit['category_id']
 
        if m_image_id==m_images_id :
            for segitem in annit['segmentation']:
                points = []
                for idx in range(0,len(segitem),2):
                    x,y=segitem[idx],segitem[idx+1]
                    if str(x).isalpha() or str(y).isalpha():
                        break
                    points.append([x,y])
                itemData = {'points': []}
                if len(points)==0:
                    flag = False
                    break

                itemData["label"] =label[m_category_id-1]
                itemData['points'].extend(points)
                itemData["group_id"] = None
                itemData["shape_type"] = "polygon"
                itemData["flag"] = {}
                data["shapes"].append(itemData)

    data['imagePath'] = m_filename
    data['imageData'] = None
    data['imageHeight'] = m_images_height
    data['imageWidth'] = m_images_width
        
    jsonName = ".".join([m_name, "json"])
    jpgName = ".".join([m_name, "jpg"])
    print(jsonName)
    jsonPath = os.path.join(labelme_json_img_path, jsonName)
    jpgPath = os.path.join(labelme_json_img_path, jpgName)
    with open(jsonPath, "w") as f:
        json.dump(data, f, indent = 2)
    cv2.imwrite(jpgPath, m_image)

print('{} files'.format(len(os.listdir(labelme_json_img_path))))
print("加载入文件完成...")

2. 根据个人需要，key_words 修改成需要类别，如person

# -*- coding: utf-8 -*-
import glob
import os
import cv2
import json
import io

"""
要将官方的coco数据集转成labelme，稍微有点不同，因为官方的coco的实际标签是排列到90，当将coco数据集转成labelme数据集时候，需要修改一下标签，插入空标签

分别要在（从１开始）　12 26 29 30 45 66 68 69 71 83 插入空labelme
"""

coco=["person", "bicycle", "car", "motorcycle", "airplane", "bus", "train", "truck", "boat", "traffic light",
        "fire hydrant","" ,"stop sign", "parking meter", "bench", "bird", "cat", "dog", "horse", "sheep", "cow",
        "elephant", "bear", "zebra", "giraffe", "","backpack", "umbrella","", "","handbag", "tie", "suitcase", "frisbee",
        "skis", "snowboard", "sports ball", "kite", "baseball bat", "baseball glove", "skateboard", "surfboard",
        "tennis racket", "bottle","", "wine glass", "cup", "fork", "knife", "spoon", "bowl", "banana", "apple",
        "sandwich", "orange", "broccoli", "carrot", "hot dog", "pizza", "donut", "cake", "chair", "couch",
        "potted plant", "bed","", "dining table", "","","toilet", "","tv", "laptop", "mouse", "remote", "keyboard", "cell phone",
        "microwave", "oven", "toaster", "sink", "refrigerator", "book", "","clock", "vase", "scissors", "teddy bear",
        "hair drier", "toothbrush"]

key_words = 'person'

label=dict()
for idx,item in enumerate(coco):
    label.update({idx:item})

labelme_json_img_path= "/home/ubuntu/pick_datasets/coco_val_person"    # 保存数据集
coco_json_path = '/home/ubuntu/pick_datasets/annotations_trainval2017/annotations'  # coco annotations
jpg_path= "/home/ubuntu/pick_datasets/val2017"  # coco val images
coco_json=glob.glob(os.path.join(coco_json_path,"*.json"))[2]
# print(coco_json)
file_json = io.open(coco_json,'r',encoding='utf-8')
m_json_data = file_json.read()
m_data = json.loads(m_json_data)
#m_type=m_data['type']
 
for item in m_data['images']:
    flag=False
    m_images_file_name = item['file_name']
    (filename_path, m_filename) = os.path.split(m_images_file_name)
    (m_name, extension) = os.path.splitext(m_filename)
    m_image=cv2.imread(os.path.join(jpg_path,m_name+".jpg"))
    m_images_height = item['height']
    m_images_width = item['width']
    m_images_id = item['id']
    data = {}
    data['version'] = "5.1.1"
    data['flags'] = {}
    data["shapes"] = []
    for annit in m_data['annotations']:
        m_image_id=annit['image_id']
        m_category_id=annit['category_id']
 
        if m_image_id==m_images_id and label[m_category_id-1]==key_words:
            flag = True
            for segitem in annit['segmentation']:
                points = []
                for idx in range(0,len(segitem),2):
                    x,y=segitem[idx],segitem[idx+1]
                    if str(x).isalpha() or str(y).isalpha():
                        break
                    points.append([x,y])
                itemData = {'points': []}
                itemData["label"] =label[m_category_id-1]
                itemData['points'].extend(points)
                itemData["group_id"] = None
                itemData["shape_type"] = "polygon"
                itemData["flag"] = {}
                data["shapes"].append(itemData)
    if flag:
        data['imagePath'] = m_filename
        data['imageData'] = None
        data['imageHeight'] = m_images_height
        data['imageWidth'] = m_images_width
            
        jsonName = ".".join([m_name, "json"])
        jpgName = ".".join([m_name, "jpg"])
        print(jsonName)
        jsonPath = os.path.join(labelme_json_img_path, jsonName)
        jpgPath = os.path.join(labelme_json_img_path, jpgName)
        with open(jsonPath, "w") as f:
            json.dump(data, f, indent = 2)
        cv2.imwrite(jpgPath, m_image)

print('{} files'.format(len(os.listdir(labelme_json_img_path))))
print("加载入文件完成...")