mask rcnn训练自己的数据集_fasterrcnn训练自己的数据集

实例分割是针对同一个类别的不同个体或者不同部分之间进行区分
我的任务是对同一个类别的不同个体进行区分，在标注的时候，不同的个体需要设置不同的标签名称

在进行标注的时候不要勾选 labelme 界面左上角 File 下拉菜单中的 Stay With Images Data 选项
否则生成的json会包含 Imagedata 信息（是很长的一大串加密的软链接），会占用很大的内存

__ignore__
__background__
seedling #根据自己的实际情况更改

分别存放的训练集和测试集图片和整合后的标签文件
seed_train 
seed_val

把整合后的标签文件剪切复制到同级目录下
seed_train_annotation.josn
seed_val_annotation.json

说明：
一次只能操作一个文件夹，也就是说：
训练集生成需要执行一次代码
测试集生成就需要更改路径之后再执行一次代码

import argparse
import collections
import datetime
import glob
import json
import os
import os.path as osp
import sys
import uuid
import time

import imgviz
import numpy as np

import labelme

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

#https://github.com/pascal1129/kaggle_airbus_ship_detection/blob/master/0_rle_to_coco/1_ships_to_coco.py

def main():
    parser = argparse.ArgumentParser(description="json2coco")
    #原始json文件保存的路径
    parser.add_argument("--input_dir", help="input annotated directory",default="E:/Deep_learning/seed-mask/data/seed/seed_train")
    #整合后的json文件保存的路径
    parser.add_argument("--output_dir", help="output dataset directory",default="E:/Deep_learning/seed-mask/data/seed/datasets/seed_train")
    parser.add_argument("--labels", help="labels file", default='E:/Deep_learning/seed-mask/data/seed/labels.txt')#required=True
    parser.add_argument( "--noviz", help="no visualization", action="store_true" ,default="--noviz")
    args = parser.parse_args()

    now = datetime.datetime.now()
    start= time.time()

    data = dict(
        info=dict(
            description="seedling datasets",
            url=None,
            version="label=4.5.6",
            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
        ],
    )

    class_name_to_id = { 
   }
    for i, line in enumerate(open(args.labels).readlines()):
        class_id = i - 1  # starts with -1
        class_name = line.strip()
        if class_id == -1:
            assert class_name == "__ignore__"
            continue
        if class_id == 0:
            assert class_name == "__background__"
            continue        
        class_name_to_id[class_name] = class_id
        #print(class_id,class_name,'\n')
        data["categories"].append(
            dict(supercategory="seedling", id=class_id, name=class_name,)#一类目标+背景，id=0表示背景
        )
    print("categories 生成完成",'\n')

    out_ann_file = osp.join(args.output_dir, "seed_train_anno.json")#自动添加"/" 这里要改 
    
    
    label_files = glob.glob(osp.join(args.input_dir, "*.json"))#图像id从json文件中读取
    for image_id, filename in enumerate(label_files):
        print(image_id, filename)
        #print("Generating dataset from:", filename)

        label_file = labelme.LabelFile(filename=filename)

        base = osp.splitext(osp.basename(filename))[0]#图片名
        out_img_file = osp.join(args.output_dir, base + ".jpg")# 保存图片路径

        img = labelme.utils.img_data_to_arr(label_file.imageData)
        imgviz.io.imsave(out_img_file, img)
        data["images"].append(
            dict(
                #license=0,
                #url=None,
                file_name=osp.relpath(out_img_file, osp.dirname(out_ann_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.shape_to_mask(img.shape[:2], points, shape_type)#labelme=4.5.6的shape_to_mask函数
            if group_id is None:
                group_id = uuid.uuid1()

            instance = (label, group_id)
            #print(instance)

            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=1,#都是1类cls_id
                    segmentation=segmentations[instance],
                    area=area,
                    bbox=bbox,
                    iscrowd=0,
                )
            )
    
    print("annotations 生成完成",'\n')

# 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:
        json.dump(data, f,indent = 2)
        
    cost_time =(time.time()-start)/1000
    print("cost_time:{:.2f}s".format(cost_time) )


if __name__ == "__main__":
    main()