复现mask2former并跑通自定义数据集

1.json数据集转coco格式

我们一般使用如labelme标注软件对数据集进行标注，保存为json文件。但这里的mask2former只支持固定数据格式，我们转为通用性很高的coco格式。通过拉取代码文件。

git clone https://gitcode.com/fcakyon/labelme2coco.git
cd labelme2coco-master
python setup.py install

注意image要和json放在同一文件夹下：

使用如下代码，将标注的json文件转成coco格式的json。

import os
import json
import PIL.Image
import PIL.ImageDraw
import numpy as np
from labelme2coco.utils import create_dir, list_jsons_recursively
from labelme2coco.image_utils import read_image_shape_as_dict


class labelme2coco(object):
    def __init__(self, labelme_folder='', save_json_path='./new.json'):
        """
        Args:
            labelme_folder: folder that contains labelme annotations and image files
            save_json_path: path for coco json to be saved
        """
        self.save_json_path = save_json_path
        self.images = []
        self.categories = []
        self.annotations = []
        self.label = []
        self.annID = 1
        self.height = 0
        self.width = 0

        # create save dir
        save_json_dir = os.path.dirname(save_json_path)
        create_dir(save_json_dir)

        # get json list
        _, labelme_json = list_jsons_recursively(labelme_folder)
        self.labelme_json = labelme_json

        self.save_json()

    def data_transfer(self):
        for num, json_path in enumerate(self.labelme_json):
            with open(json_path, 'r') as fp:
                # load json
                data = json.load(fp)
#                (prefix, res) = os.path.split(json_path)
#                (file_name, extension) = os.path.splitext(res)
                self.images.append(self.image(data, num, json_path))
                for shapes in data['shapes']:
                    label = shapes['label']
                    if label not in self.label:
                        self.categories.append(self.category(label))
                        self.label.append(label)
                    points = shapes['points']
                    self.annotations.append(self.annotation(points, label, num))
                    self.annID += 1

    def image(self, data, num, json_path):
        image = {}
        # get image path
        _, img_extension = os.path.splitext(data["imagePath"])
        image_path = json_path.replace(".json", img_extension)
        img_shape = read_image_shape_as_dict(image_path)
        height, width = img_shape['height'], img_shape['width']

        image['height'] = height
        image['width'] = width
        image['id'] = int(num + 1)
        image['file_name'] = image_path

        self.height = height
        self.width = width

        return image

    def category(self, label):
        category = {}
        category['supercategory'] = label
        category['id'] = int(len(self.label) + 1)
        category['name'] = label

        return category

    def annotation(self, points, label, num):
        annotation = {}
        annotation['iscrowd'] = 0
        annotation['image_id'] = int(num + 1)

        annotation['bbox'] = list(map(float, self.getbbox(points)))

        # coarsely from bbox to segmentation
        x = annotation['bbox'][0]
        y = annotation['bbox'][1]
        w = annotation['bbox'][2]
        h = annotation['bbox'][3]
        annotation['segmentation'] = [np.asarray(points).flatten().tolist()]

        annotation['category_id'] = self.getcatid(label)
        annotation['id'] = int(self.annID)
        # add area info
        annotation['area'] = self.height * self.width  # the area is not used for detection
        return annotation

    def getcatid(self, label):
        for categorie in self.categories:
            if label == categorie['name']:
                return categorie['id']
            # if label[1]==categorie['name']:
            #     return categorie['id']
        return -1

    def getbbox(self, points):
        # img = np.zeros([self.height,self.width],np.uint8)
        # cv2.polylines(img, [np.asarray(points)], True, 1, lineType=cv2.LINE_AA)
        # cv2.fillPoly(img, [np.asarray(points)], 1)
        polygons = points
        mask = self.polygons_to_mask([self.height, self.width], polygons)
        return self.mask2box(mask)

    def mask2box(self, mask):
        # np.where(mask==1)
        index = np.argwhere(mask == 1)
        rows = index[:, 0]
        clos = index[:, 1]

        left_top_r = np.min(rows)  # y
        left_top_c = np.min(clos)  # x

        right_bottom_r = np.max(rows)
        right_bottom_c = np.max(clos)

        return [left_top_c, left_top_r, right_bottom_c-left_top_c, right_bottom_r-left_top_r]  # [x1,y1,w,h] for coco box format

    def polygons_to_mask(self, img_shape, polygons):
        mask = np.zeros(img_shape, dtype=np.uint8)
        mask = PIL.Image.fromarray(mask)
        xy = list(map(tuple, polygons))
        PIL.ImageDraw.Draw(mask).polygon(xy=xy, outline=1, fill=1)
        mask = np.array(mask, dtype=bool)
        return mask

    def data2coco(self):
        data_coco = {}
        data_coco['images'] = self.images
        data_coco['categories'] = self.categories
        data_coco['annotations'] = self.annotations
        return data_coco

    def save_json(self):
        self.data_transfer()
        self.data_coco = self.data2coco()

        json.dump(self.data_coco, open(self.save_json_path, 'w', encoding='utf-8'), indent=4, separators=(',', ': '), cls=MyEncoder)


# type check when save json files
class MyEncoder(json.JSONEncoder):
    def default(self, obj):
        if isinstance(obj, np.integer):
            return int(obj)
        elif isinstance(obj, np.floating):
            return float(obj)
        elif isinstance(obj, np.ndarray):
            return obj.tolist()
        else:
            return super(MyEncoder, self).default(obj)


if __name__ == "__main__":
    #labelme_folder 你的标注图片和标签所在的文件夹
    labelme_folder = r"D:\jiedan\tool_data\hh\val"
    #save_json_path 转换生成的coco格式的标签文件的保存路径
    save_json_path = r"D:\jiedan\tool_data\hh\val_coco_format.json"
    labelme2coco(labelme_folder, save_json_path)
	#下面是可视化标注的mask
    # import os
    #
    # from pycocotools.coco import COCO
    # from skimage import io
    # from matplotlib import pyplot as plt
    #
    # json_file = r'D:\jiedan\dataset\label\train_coco_format.json'  # 输入文件路径
    # dataset_dir = r''
    # coco = COCO(json_file)
    # catIds = coco.getCatIds(catNms=['head'])  # 标注的图片的不同类型别，超过一类，用逗号隔开
    # imgIds = coco.getImgIds(catIds=catIds)  # 图片id，许多值
    # for i in range(len(imgIds)):
    #     img = coco.loadImgs(imgIds[i])[0]
    #     I = io.imread(dataset_dir + img['file_name'])
    #     plt.axis('off')
    #     plt.imshow(I)  # 绘制图像，显示交给plt.show()处理
    #     annIds = coco.getAnnIds(imgIds=img['id'], catIds=catIds, iscrowd=None)
    #     anns = coco.loadAnns(annIds)
    #     coco.showAnns(anns)
    #     plt.show()  # 显示图像
    #
    #

2.注册mask2former自定数据集

转换好的数据集存放位置如下：
train和val文件夹存放划分好的数据集图片，两个json文件对应于生成的json转好的coco格式。

我们在主函数中进行数据集注册：

// 第1个参数为自定义数据集名称
//第2个参数为json文件的相对路径
//第3个参数为图片数据集文件的相对路径
from detectron2.data.datasets import register_coco_instances
register_coco_instances("my_dataset_train", {}, "tool_data/train_coco_format.json", "tool_data/train")
register_coco_instances("my_dataset_val", {}, "tool_data/val_coco_format.json", "tool_data/val")

3.更改配置文件

我们是使用的coco格式的数据集进行训练，所以找到配置文件中的coco文件

首先我们对基础base配置文件进行更改，将train和test数据集换成我们刚在在主函数中注册的数据集名称。

接下来，比如我们使用maskformer2_R50_bs16_50ep.yaml模型文件，则进行相应的类别数进行修改。

最后就可以训练自己的数据集啦！（ps：模型的权重去model ZOO中进行下载即可！）

python train_net.py --num-gpus 1 --config-file configs/coco/instance-segmentation/maskformer2_R50_bs16_50ep.yaml  MODEL.WEIGHTS "weights/model_final_94dc52.pkl"

4.小工具

最后再附上一个通过标注的json文件转成mask的代码：

import os
import json
import base64
import imgviz
import PIL.Image
import os.path as osp

from tqdm import tqdm
from labelme import utils
from threading import Thread
'''这个是根据labelme标注的json文件生成mask'''

def ConvertOne(labelme_dir, json_file, save_dir, label_name_to_value):
    out_dir = os.path.join(save_dir, json_file.replace(".json", ""))
    if not os.path.exists(out_dir):
        os.makedirs(out_dir)
    json_path = osp.join(labelme_dir, json_file)
    with open(json_path, "r") as jf:
        data = json.load(jf)
        imageData = data.get("imageData")

        # labelme 的图像数据编码以及返回处理格式
        if not imageData:
            imagePath = os.path.join(os.path.dirname(json_file), data["imagePath"])
            with open(imagePath, "rb") as f:
                imageData = f.read()
                imageData = base64.b64encode(imageData).decode("utf-8")
        img = utils.img_b64_to_arr(imageData)

        lbl, _ = utils.shapes_to_label(
            img.shape, data["shapes"], label_name_to_value
        )

        label_names = [None] * (max(label_name_to_value.values()) + 1)

        for name, value in label_name_to_value.items():
            label_names[value] = name

        # label_names={'_background_','line','border'}
        lbl_viz = imgviz.label2rgb(
            lbl, imgviz.asgray(img), label_names=label_names, loc="rb"
        )

        PIL.Image.fromarray(img).save(osp.join(out_dir, "img.png"))
        # 保存标签图片
        utils.lblsave(osp.join(out_dir, "label.png"), lbl)
        # utils.lblsave((out_dir + ".png"), lbl)
        # 保存带标签的可视化图像
        PIL.Image.fromarray(lbl_viz).save(osp.join(out_dir, "label_viz.png"))

        with open(osp.join(out_dir, "label_names.txt"), "w") as f:
            for lbl_name in label_names:
                f.write(lbl_name + "\n")


def main():
    labelme_dir = r'D:\jiedan\tool_data\data_annotation'  # json文件存放文件夹
    save_dir = r'D:\jiedan\tool_data\output'  # 结果生成文件夹
    # 类别标签
    class_names = {
        '_background_': 0,
        "tool 1": 1,
        "tool 2": 2,
        "tool 3": 3,
        "tool 4": 4,
        "tool 5": 5,
        "tool 6": 6,
        "tool 7": 7,
        "tool 8": 8,

    }
    # 列出labelme勾画标签后文件夹中保存的所有文件名
    file_list = os.listdir(labelme_dir)
    # 找到勾画保存的所有json标签
    json_list = []
    [json_list.append(x) for x in file_list if x.endswith(".json")]

    for json_file in tqdm(json_list):
        # 单线程
        ConvertOne(labelme_dir, json_file, save_dir, class_names)

        # 多线程
        # Thread(target=ConvertOne, args=(labelme_dir, json_file, save_dir, class_names)).start()
        print(f"生成结果保存地址：{save_dir}")


if __name__ == "__main__":
    main()

生成结果为：