labelme轉coco數據集

原始labelme數據目錄結構如下：

|-- images
|     |---  1.jpg
|     |---  1.json
|     |---  2.jpg
|     |---  2.json
|     |---  .......
|-- labelme2coco.py
|-- labels.txt

• imges目錄下就是你的數據集原始圖片，加上labelme標注的json文件。

• labelme2coco.py源碼放到最后。

• labels.txt就是你的類別標簽，假設我有兩個類（lm，ls），那么對應的labels.txt內容如下：

__ignore__
_background_
lm
ls

在labelme2coco.py文件的目錄下，打開命令行執行：

 python labelme2coco.py --input_dir images --output_dir coco --labels labels.txt

• --input_dir：指定images文件夾

• --output_dir：指定你的輸出文件夾

• --labels：指定你的labels.txt文件

執行結果如下圖：

生成的coco數據集目錄結構如下：

|-- annotations
| 	|---  instances_train2017.json
|       |---  instances_val2017.json
|-- train2017
| 	|---  2.jpg
| 	|---  5.jpg
| 	|---  .......
|-- val2017
| 	|---  1.jpg
| 	|---  3.jpg
| 	|---  .......
|-- visualization
| 	|---  1.jpg
| 	|---  2.jpg
| 	|---  .......

訓練之用前三個文件夾就可以了，也就是annotations,train2017,val2017就可以了。visualization可以用來觀察自己標注的數據集效果。源碼里都標注好了，大部分來自labelme官方的源碼：labelme/labelme2coco.py · GitHub

如果想調整訓練集驗證集的比例，可以在labelme2coco.py源碼中搜索 test_size

labelme2coco.py源碼：

# 命令行執行：  python labelme2coco.py --input_dir images --output_dir coco --labels labels.txt
# 輸出文件夾必須為空文件夾

import argparse
import collections
import datetime
import glob
import json
import os
import os.path as osp
import sys
import uuid
import imgviz
import numpy as np
import labelme
from sklearn.model_selection import train_test_split

try:
    import pycocotools.mask
except ImportError:
    print("Please install pycocotools:\n\n    pip install pycocotools\n")
    sys.exit(1)


def to_coco(args,label_files,train):

    # 創建 總標簽data 
    now = datetime.datetime.now()
    data = dict(
        info=dict(
            description=None,
            url=None,
            version=None,
            year=now.year,
            contributor=None,
            date_created=now.strftime("%Y-%m-%d %H:%M:%S.%f"),
        ),
        licenses=[dict(url=None, id=0, name=None,)],
        images=[
            # license, url, file_name, height, width, date_captured, id
        ],
        type="instances",
        annotations=[
            # segmentation, area, iscrowd, image_id, bbox, category_id, id
        ],
        categories=[
            # supercategory, id, name
        ],
    )

    # 創建一個 {類名 : id} 的字典，並保存到 總標簽data 字典中。
    class_name_to_id = {}
    for i, line in enumerate(open(args.labels).readlines()):
        class_id = i - 1  # starts with -1
        class_name = line.strip()   # strip() 方法用於移除字符串頭尾指定的字符（默認為空格或換行符）或字符序列。
        if class_id == -1:
            assert class_name == "__ignore__"   # background:0, class1:1, ,,
            continue
        class_name_to_id[class_name] = class_id
        data["categories"].append(
            dict(supercategory=None, id=class_id, name=class_name,)
        )


    if train:
        out_ann_file = osp.join(args.output_dir, "annotations","instances_train2017.json")
    else:
        out_ann_file = osp.join(args.output_dir, "annotations","instances_val2017.json")


    for image_id, filename in enumerate(label_files):

        label_file = labelme.LabelFile(filename=filename)
        base = osp.splitext(osp.basename(filename))[0]      # 文件名不帶后綴
        if train:
            out_img_file = osp.join(args.output_dir, "train2017", base + ".jpg")
        else:
            out_img_file = osp.join(args.output_dir, "val2017", base + ".jpg")
        
        print("| ",out_img_file)


        # ************************** 對圖片的處理開始 *******************************************
        # 將標簽文件對應的圖片進行保存到對應的 文件夾。train保存到 train2017/ test保存到 val2017/
        img = labelme.utils.img_data_to_arr(label_file.imageData)   # .json文件中包含圖像，用函數提出來
        imgviz.io.imsave(out_img_file, img)     # 將圖像保存到輸出路徑

        # ************************** 對圖片的處理結束 *******************************************

        # ************************** 對標簽的處理開始 *******************************************
        data["images"].append(
            dict(
                license=0,
                url=None,
                file_name=base+".jpg",              # 只存圖片的文件名
                # file_name=osp.relpath(out_img_file, osp.dirname(out_ann_file)),  # 存標簽文件所在目錄下找圖片的相對路徑

                ##   out_img_file : "/coco/train2017/1.jpg"
                ##   out_ann_file : "/coco/annotations/annotations_train2017.json"
                ##   osp.dirname(out_ann_file) : "/coco/annotations"
                ##   file_name : ..\train2017\1.jpg   out_ann_file文件所在目錄下 找 out_img_file 的相對路徑
                height=img.shape[0],
                width=img.shape[1],
                date_captured=None,
                id=image_id,
            )
        )

        masks = {}  # for area
        segmentations = collections.defaultdict(list)  # for segmentation
        for shape in label_file.shapes:
            points = shape["points"]
            label = shape["label"]
            group_id = shape.get("group_id")
            shape_type = shape.get("shape_type", "polygon")
            mask = labelme.utils.shape_to_mask(
                img.shape[:2], points, shape_type
            )

            if group_id is None:
                group_id = uuid.uuid1()

            instance = (label, group_id)

            if instance in masks:
                masks[instance] = masks[instance] | mask
            else:
                masks[instance] = mask

            if shape_type == "rectangle":
                (x1, y1), (x2, y2) = points
                x1, x2 = sorted([x1, x2])
                y1, y2 = sorted([y1, y2])
                points = [x1, y1, x2, y1, x2, y2, x1, y2]
            else:
                points = np.asarray(points).flatten().tolist()

            segmentations[instance].append(points)
        segmentations = dict(segmentations)

        for instance, mask in masks.items():
            cls_name, group_id = instance
            if cls_name not in class_name_to_id:
                continue
            cls_id = class_name_to_id[cls_name]

            mask = np.asfortranarray(mask.astype(np.uint8))
            mask = pycocotools.mask.encode(mask)
            area = float(pycocotools.mask.area(mask))
            bbox = pycocotools.mask.toBbox(mask).flatten().tolist()

            data["annotations"].append(
                dict(
                    id=len(data["annotations"]),
                    image_id=image_id,
                    category_id=cls_id,
                    segmentation=segmentations[instance],
                    area=area,
                    bbox=bbox,
                    iscrowd=0,
                )
            )
        # ************************** 對標簽的處理結束 *******************************************

        # ************************** 可視化的處理開始 *******************************************
        if not args.noviz:
            labels, captions, masks = zip(
                *[
                    (class_name_to_id[cnm], cnm, msk)
                    for (cnm, gid), msk in masks.items()
                    if cnm in class_name_to_id
                ]
            )
            viz = imgviz.instances2rgb(
                image=img,
                labels=labels,
                masks=masks,
                captions=captions,
                font_size=15,
                line_width=2,
            )
            out_viz_file = osp.join(
                args.output_dir, "visualization", base + ".jpg"
            )
            imgviz.io.imsave(out_viz_file, viz)
        # ************************** 可視化的處理結束 *******************************************

    with open(out_ann_file, "w") as f:  # 將每個標簽文件匯總成data后，保存總標簽data文件
        json.dump(data, f)


# 主程序執行
def main():
    parser = argparse.ArgumentParser(
        formatter_class=argparse.ArgumentDefaultsHelpFormatter
    )
    parser.add_argument("--input_dir", help="input annotated directory")
    parser.add_argument("--output_dir", help="output dataset directory")
    parser.add_argument("--labels", help="labels file", required=True)
    parser.add_argument("--noviz", help="no visualization", action="store_true")
    args = parser.parse_args()

    if osp.exists(args.output_dir):
        print("Output directory already exists:", args.output_dir)
        sys.exit(1)
    os.makedirs(args.output_dir)
    print("| Creating dataset dir:", args.output_dir)
    if not args.noviz:
        os.makedirs(osp.join(args.output_dir, "visualization"))

    # 創建保存的文件夾
    if not os.path.exists(osp.join(args.output_dir, "annotations")):
        os.makedirs(osp.join(args.output_dir, "annotations"))
    if not os.path.exists(osp.join(args.output_dir, "train2017")):
        os.makedirs(osp.join(args.output_dir, "train2017"))
    if not os.path.exists(osp.join(args.output_dir, "val2017")):
        os.makedirs(osp.join(args.output_dir, "val2017"))

    # 獲取目錄下所有的.jpg文件列表
    feature_files = glob.glob(osp.join(args.input_dir, "*.jpg"))
    print('| Image number: ', len(feature_files))

    # 獲取目錄下所有的joson文件列表
    label_files = glob.glob(osp.join(args.input_dir, "*.json"))
    print('| Json number: ', len(label_files))
    

    # feature_files:待划分的樣本特征集合    label_files:待划分的樣本標簽集合    test_size:測試集所占比例 
    # x_train:划分出的訓練集特征      x_test:划分出的測試集特征     y_train:划分出的訓練集標簽    y_test:划分出的測試集標簽
    x_train, x_test, y_train, y_test = train_test_split(feature_files, label_files, test_size=0.3)
    print("| Train number:", len(y_train), '\t Value number:', len(y_test))

    # 把訓練集標簽轉化為COCO的格式，並將標簽對應的圖片保存到目錄 /train2017/
    print("—"*50) 
    print("| Train images:")
    to_coco(args,y_train,train=True)
    
    # 把測試集標簽轉化為COCO的格式，並將標簽對應的圖片保存到目錄 /val2017/ 
    print("—"*50)
    print("| Test images:")
    to_coco(args,y_test,train=False)
    

if __name__ == "__main__":
    print("—"*50)
    main()
    print("—"*50)