mask2former训练自定义数据集

一 、下载工程源码
项目工程源码

二、环境配置
根据实际需要可能需要执行TORCH_CUDA_ARCH_LIST='11.6' FORCE_CUDA=1 python setup.py build install
1、 下载所需的detectron2 或者使用git clone git@github.com:facebookresearch/detectron2.git获取
2、将detectron2放到Mask2Former根目录下
3、pip install -e .
4、pip install git+https://github.com/cocodataset/panopticapi.git
5、pip install git+https://github.com/mcordts/cityscapesScripts.git
6、然后执行项目自带推荐环境requirements.txt
7、cd mask2former/modeling/pixel_decoder/ops/
8、sh make.sh

三、准备数据集
在项目根目录下创建数据集文件夹，例如

datas
	|__train
		|__1.pg
		|__1.json
		|__2.jpg
		|__2.jsom
		|__3.jpg
		|__3.json
	|__val
		|__1.pg
		|__1.json

将标注好的数据进行转换
json2coco.py

import labelme2coco

# 设置LabelMe注解和图片文件所在的目录
labelme_folder = "datas/train"					# 修改为对应路径
# 设置导出文件的目录
export_dir = "datas"										# 修改为对应路径

# 设置训练集和验证集的切分比例
train_split_rate = 0.95
# 设置类别ID的起始值
category_id_start = 1

# 转换LabelMe注解到COCO格式
labelme2coco.convert(
    labelme_folder,
    export_dir,
    train_split_rate=train_split_rate,
    category_id_start=category_id_start
)

执行完后，会在hh文件夹下生成train_coco_format.json、val_coco_format.json

datas
	|__train
		|__1.pg
		|__1.json
		|__2.jpg
		|__2.jsom
		|__3.jpg
		|__3.json
	|__val
		|__1.pg
		|__1.json
	|__train_coco_format.json
	|__val_coco_format.json

四、注册数据集，添加

register_coco_instances("my_dataset_train", {}, "datas/train_coco_format.json", "")
register_coco_instances("my_dataset_val", {}, "datas/val_coco_format.json", "")

添加后如下

if __name__ == "__main__":
    args = default_argument_parser().parse_args()
    print("Command Line Args:", args)

    register_coco_instances("my_dataset_train", {}, "datas/train_coco_format.json", "")
	register_coco_instances("my_dataset_val", {}, "datas/val_coco_format.json", "")

    launch(
        main,
        args.num_gpus,
        num_machines=args.num_machines,
        machine_rank=args.machine_rank,
        dist_url=args.dist_url,
        args=(args,),
    )

五、更改配置文件
修改configs/coco/instance-segmentatio/Base-COCO-InstanceSegmentation.yaml中修改为注册的数据集名。

DATASETS:
  TRAIN: ("my_dataset_train",)
  TEST: ("my_dataset_val",)

修改configs/coco/instance-segmentatio/maskformer2_R50_bs16_50ep.yaml中修改为分类识别个数。

NUM_CLASSES: 2				# 修改为识别类别个数

六、训练

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

七、预测

python demo/demo.py  --config-file ../configs/coco/instance-segmentation/maskformer2_R50_bs16_50ep.yaml --input "../datas/val/*.jpg" --confidence-threshold 0.3 --output "../output/test_mask/" --opts MODEL.WEIGHTS ../output/model_0024999.pth