mask2former训练自己的分割数据集（包含遇见的问题及解决办法）

一、环境配置

1.1 下载所需源码

mask2former:

https://github.com/facebookresearch/Mask2Former/tree/main

detectron2:

https://github.com/facebookresearch/detectron2

解压后，将两个文件夹放置同一目录。

1.2 配置环境

        1）anaconda下新建并激活虚拟环境

新建：
conda create -n mask2former python==3.8
激活：
source activate mask2former

        2）安装torch、torchvision

                在如下路径中，找对应的版本下载：

https://download.pytorch.org/whl/cu113

                我的是cu113，所需版本为：

                torch-1.12.1+cu113-cp38-cp38-linux_x86_64.whl

                torchvision-0.13.1+cu113-cp38-cp38-linux_x86_64.whl

                之后，运行如下命令即可：

pip install -i https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple torch--***.whl

        3）安装detectron2

                cd到detectron2所在的根目录下，不要 cd 进detectron2文件夹：

python -m pip install -e detectron2

        4）安装mask2former所需环境和MSDA模块

                cd进mask2former文件夹，运行：

pip install -r requirements.txt

                然后，cd 进 mask2former/modeling/pixel_decoder/ops/目录下，运行：

sh make.sh

1.3 配置环境容易报错之处：

 1）MSDA编译时报错（上一步sh make.sh操作）：

原因：setuptools版本过高，我安装的是setuptools=70，卸载安装为60版本即可：

pip install setuptools==60

        2）detectron2遇到：detectron2/_C.cpython-38-x86_64-linux-gnu.so: undefined symbol: _ZNK2at6Tensor7is_cudaEv

原因：更新torch版本后，Detectron2需要重新安装.

cd到detectron2所在目录：

python -m pip install -e detectron2

二、配置数据集

在项目根目录下创建数据集文件夹，例如：

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

json2coco的转换：

json2coco.py：

import labelme2coco

# 设置LabelMe的json标注文件和图片文件所在的目录
labelme_folder = "/home/houxiangguan/HPA/hpa_train_val_data/mask2former/train"					# 修改为对应路径
# 设置导出文件的目录
export_dir = "/home/houxiangguan/HPA/hpa_train_val_data/mask2former"										# 修改为对应路径

# 设置训练集和验证集的切分比例
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。

三、修改代码

3.1 注册数据集

修改train_net.py如下：

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,),
    )

3.2 修改配置文件

修改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				# 修改为识别类别个数

3.3 训练并预测

训练：

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