利用maskFormer/mask2Former配置环境并训练自己的数据集

原创

已于 2023-12-09 20:25:08 修改 · 9.6k 阅读

67 ·

CC 4.0 BY-SA版权

文章标签：

#深度学习 #人工智能 #python #pytorch

于 2023-12-08 19:15:45 首次发布

1. 配置环境:

使用detectron2配置的环境，所以也是复制之前电脑上用过detectron2的环境

conda create -n mask2former --clone PatchDCT
conda activate mask2former

然后在requirement里面下载需要的包

另外在mask2former中还需要make一下：

2. 数据集配置

利用之前配置好的coco数据集，但由于maskFormer/mask2Former使用的数据集是包含实例分割和语义分割，还需要在datasets/prepare_ade20k_sem_seg.py输出符合代码运行的数据集格式。

环境里面还需要导入包：

pip install git+https://github.com/cocodataset/panopticapi.git

在trin_net.py加入：（这个在之前创建自己的coco数据集也有提到）

# 自己的数据集
from detectron2.data import DatasetCatalog
from detectron2.data.datasets.coco import load_coco_json
os.environ['CUDA_LAUNCH_BLOCKING'] = '1'
# 声明类别，尽量保持  #需修改成为自己的类别名称
# CLASS_NAMES = ["_background_", "person", "bicycle", "car", "UAV"] # 没有背景就不标注背景
CLASS_NAMES = ["person", "bicycle", "car", "UAV"]
# 数据集类别元数据
DATASET_CATEGORIES = [
    {"name": "person", "id": 0, "isthing": 1, "color": [219, 142, 185]},
    {"name": "bicycle", "id": 1, "isthing": 1, "color": [119, 11, 32]},
    {"name": "car", "id": 2, "isthing": 1, "color": [0, 0, 142]},
    {"name": "UAV", "id": 3, "isthing": 1, "color": [119, 11, 32]},
]

# 数据集路径 此路径需修改为自己的路径
DATASET_ROOT = r'G:/coco_4classes'
ANN_ROOT = os.path.join(DATASET_ROOT, 'annotations')

TRAIN_PATH = os.path.join(DATASET_ROOT, 'train2017')
VAL_PATH = os.path.join(DATASET_ROOT, 'val2017')

TRAIN_JSON = os.path.join(ANN_ROOT, 'instances_train2017.json')
# VAL_JSON = os.path.join(ANN_ROOT, 'val.json')
VAL_JSON = os.path.join(ANN_ROOT, 'instances_val2017.json')
# print(TRAIN_PATH)
# 声明数据集的子集 coco_my_train  修改为自己的数据集名字
PREDEFINED_SPLITS_DATASET = {
    "coco_my_train": (TRAIN_PATH, TRAIN_JSON),
    "coco_my_val": (VAL_PATH, VAL_JSON),
}


# 注册数据集
def register_dataset():
    """
    purpose: register all splits of dataset with PREDEFINED_SPLITS_DATASET
    """
    for key, (image_root, json_file) in PREDEFINED_SPLITS_DATASET.items():
        register_dataset_instances(name=key,
                                   metadate=get_dataset_instances_meta(),
                                   json_file=json_file,
                                   image_root=image_root)