修改labelme原始碼,解決粘連mask分離問題

dota2職業選手發表於2020-10-10

問題:

    我們需要訓練一個粘連物體分割的模型,使用labelme進行資料標註。我在標註完以後使用labelme_json_to_dataset命令進行生成資料集的時候出現以下情況。(其中粘連部位會被其中一個類別覆蓋)

                                                 

解決方案:

    首先找到json_to_dataset.py指令碼(我的是mac系統位置是/Library/Frameworks/Python.framework/Versions/3.6/lib/python3.6/site-packages/labelme/cli/json_to_dataset.py)修改其中的以下部分:

    for shape in sorted(data["shapes"], key=lambda x: x["label"]):
        label_name = shape["label"]
        if label_name in label_name_to_value:
            label_value = label_name_to_value[label_name]
        else:
            label_value = len(label_name_to_value)
            label_name_to_value[label_name] = label_value
    lbl_0, lbl_1, _ = utils.shapes_to_label(
        img.shape, data["shapes"], label_name_to_value
    )

    然後修改shape.py指令碼中的shapes_to_label方法。(位置在/Library/Frameworks/Python.framework/Versions/3.6/lib/python3.6/site-packages/labelme/utils/shape.py)

def shapes_to_label(img_shape, shapes, label_name_to_value):
    res = []
    for shape in shapes:
        cls = np.zeros(img_shape[:2], dtype=np.int32)
        ins = np.zeros_like(cls)
        instances = []
        points = shape["points"]
        label = shape["label"]
        group_id = shape.get("group_id")
        if group_id is None:
            group_id = uuid.uuid1()
        shape_type = shape.get("shape_type", None)

        cls_name = label
        instance = (cls_name, group_id)

        if instance not in instances:
            instances.append(instance)
        ins_id = instances.index(instance) + 1
        cls_id = label_name_to_value[cls_name]

        mask = shape_to_mask(img_shape[:2], points, shape_type)
        cls[mask] = cls_id
        ins[mask] = ins_id
        res.append(cls)

    return res[0], res[1], ins

    

然後使用labelme_json_to_dataset命令再生成標籤的時候會出現以下兩張圖:

                                                          

 

相關文章