1. 程式人生 > 其它 >detectron2訓練自己的資料集及demo測試

detectron2訓練自己的資料集及demo測試

技術標籤:目標檢測深度學習

今天就寫寫detectron2訓練自己的資料。

1.環境準備

python=3.7.0
cuda=10.2.89
cudnn=7.6.5
torch=1.6.0
torchvision=0.7.0

2.資料準備(資料轉換)

1.按照yolov3的檔案目錄(相關指令碼下載提取碼tfcj):

VOC2007
----Annotations
----JPEGImages
----ImageSets
--------Main
----test.py #指令碼
----train_JPEGImages #自己建立
----val_JPEGImages #自己建立
----val_annotations #自己建立

----train_annotations #自己建立
instances_train2017.json #指令碼生成
instances_val2017.json #指令碼生成
v2c_1.py #指令碼
v2c_2.py #指令碼

首先,我們執行test.py在Main下面生成對應的.txt
2.執行v2c_1.py
我們改變裡面的引數分別將JPEGImages分為train_JPEGImages和val_JPEGImages,將Annotations分為train_annotations和val_annotations
3.執行v2c_2.py(VOC2COCO)
因為Detectron2需要的是coco資料格式,我們在這一步將VOC

格式的資料轉換為COCO資料。通過執行v2c_2.py來生成instances_train2017.jsoninstances_val017.json
4.資料準備
建立coco資料夾,並將train_JPEGImagesval_JPEGImages放入裡面並改名為train2017val2017,建立annotations資料夾並將instances_train2017.jsoninstances_val017.json(或者通過改變指令碼路徑直接生成,不需要拷貝或者改命)。其最後的目錄為:

coco
----train2017
--------xxx.jpg
----val2017
--------xxx.jpg

----annotations
--------instances_train2017.json
--------instances_val017.json

我們將coco放入detectron2/datasets/(這一步在demo測試成功後做)
最後的檔案目錄為:

3.原始碼下載編譯及demo測試

1.原始碼下載
按照官網INSTALL.md的要求按照好所需的環境。

python -m pip install 'git+https://github.com/facebookresearch/detectron2.git'
# (add --user if you don't have permission)

# Or, to install it from a local clone:
git clone https://github.com/facebookresearch/detectron2.git
python -m pip install -e detectron2
cd detectron2

其中要求gcc & g++ ≥ 5在不滿足的條件寫可執行:

sudo yum install centos-release-scl
sudo yum install devtoolset-8-gcc*
scl enable devtoolset-8 bash
gcc -v #檢視版本,只對本次會話有效

其中INSTALL.md也包括常見問題的解決方案。
2.demo測試
我們先下載預訓練模型下載存放在對應資料夾。

cd demo
python demo.py --config-file ../configs/COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml \
  --input inputpath.png --output outpath.png \
  --opts MODEL.WEIGHTS model_path.pth 

如果outpath.png顯示出對應的類別和位置則說明測試成功。

4.訓練自己的資料集(本文使用detectron2/detectron2/configconfigs/COCO-Detection/retinanet_R_50_FPN_3x.yaml)

1.註冊資料集,修改detectron2/tools/train_net.py

import logging
import os
from collections import OrderedDict
import torch

import detectron2.utils.comm as comm
from detectron2.checkpoint import DetectionCheckpointer
from detectron2.config import get_cfg
from detectron2.data import MetadataCatalog
from detectron2.engine import DefaultTrainer, default_argument_parser, default_setup, hooks, launch
from detectron2.evaluation import (
    CityscapesInstanceEvaluator,
    CityscapesSemSegEvaluator,
    COCOEvaluator,
    COCOPanopticEvaluator,
    DatasetEvaluators,
    LVISEvaluator,
    PascalVOCDetectionEvaluator,
    SemSegEvaluator,
    verify_results,
)
from detectron2.modeling import GeneralizedRCNNWithTTA

# 註冊資料集

from detectron2.data import DatasetCatalog, MetadataCatalog
from detectron2.data.datasets.coco import load_coco_json
import pycocotools
#宣告類別,儘量保持
CLASS_NAMES =["person","car","dog","car"]
# 資料集路徑
DATASET_ROOT = './datasets/coco'
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')

# 宣告資料集的子集
PREDEFINED_SPLITS_DATASET = {
    "coco_my_train": (TRAIN_PATH, TRAIN_JSON),
    "coco_my_val": (VAL_PATH, VAL_JSON),
}
#===========以下有兩種註冊資料集的方法,本人直接用的第二個plain_register_dataset的方式 也可以用register_dataset的形式==================
#註冊資料集(這一步就是將自定義資料集註冊進Detectron2)
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,
                                   json_file=json_file,
                                   image_root=image_root)


#註冊資料集例項,載入資料集中的物件例項
def register_dataset_instances(name, json_file, image_root):
    """
    purpose: register dataset to DatasetCatalog,
             register metadata to MetadataCatalog and set attribute
    """
    DatasetCatalog.register(name, lambda: load_coco_json(json_file, image_root, name))
    MetadataCatalog.get(name).set(json_file=json_file,
                                  image_root=image_root,
                                  evaluator_type="coco")

#=============================
# 註冊資料集和元資料
def plain_register_dataset():
    #訓練集
    DatasetCatalog.register("coco_my_train", lambda: load_coco_json(TRAIN_JSON, TRAIN_PATH))
    MetadataCatalog.get("coco_my_train").set(thing_classes=CLASS_NAMES,  # 可以選擇開啟,但是不能顯示中文,這裡需要注意,中文的話最好關閉
                                                    evaluator_type='coco', # 指定評估方式
                                                    json_file=TRAIN_JSON,
                                                    image_root=TRAIN_PATH)

    #DatasetCatalog.register("coco_my_val", lambda: load_coco_json(VAL_JSON, VAL_PATH, "coco_2017_val"))
    #驗證/測試集
    DatasetCatalog.register("coco_my_val", lambda: load_coco_json(VAL_JSON, VAL_PATH))
    MetadataCatalog.get("coco_my_val").set(thing_classes=CLASS_NAMES, # 可以選擇開啟,但是不能顯示中文,這裡需要注意,中文的話最好關閉
                                                evaluator_type='coco', # 指定評估方式
                                                json_file=VAL_JSON,
                                                image_root=VAL_PATH)
# 檢視資料集標註,視覺化檢查資料集標註是否正確,
#這個也可以自己寫指令碼判斷,其實就是判斷標註框是否超越影象邊界
#可選擇使用此方法
def checkout_dataset_annotation(name="coco_my_val"):
    #dataset_dicts = load_coco_json(TRAIN_JSON, TRAIN_PATH, name)
    dataset_dicts = load_coco_json(TRAIN_JSON, TRAIN_PATH)
    print(len(dataset_dicts))
    for i, d in enumerate(dataset_dicts,0):
        #print(d)
        img = cv2.imread(d["file_name"])
        visualizer = Visualizer(img[:, :, ::-1], metadata=MetadataCatalog.get(name), scale=1.5)
        vis = visualizer.draw_dataset_dict(d)
        #cv2.imshow('show', vis.get_image()[:, :, ::-1])
        cv2.imwrite('out/'+str(i) + '.jpg',vis.get_image()[:, :, ::-1])
        #cv2.waitKey(0)
        # if i == 200:
        #     break
class Trainer(DefaultTrainer):
    """
    We use the "DefaultTrainer" which contains pre-defined default logic for
    standard training workflow. They may not work for you, especially if you
    are working on a new research project. In that case you can write your
    own training loop. You can use "tools/plain_train_net.py" as an example.
    """

    @classmethod
    def build_evaluator(cls, cfg, dataset_name, output_folder=None):
        """
        Create evaluator(s) for a given dataset.
        This uses the special metadata "evaluator_type" associated with each builtin dataset.
        For your own dataset, you can simply create an evaluator manually in your
        script and do not have to worry about the hacky if-else logic here.
        """
        if output_folder is None:
            output_folder = os.path.join(cfg.OUTPUT_DIR, "inference")
        evaluator_list = []
        evaluator_type = MetadataCatalog.get(dataset_name).evaluator_type
        if evaluator_type in ["sem_seg", "coco_panoptic_seg"]:
            evaluator_list.append(
                SemSegEvaluator(
                    dataset_name,
                    distributed=True,
                    output_dir=output_folder,
                )
            )
        if evaluator_type in ["coco", "coco_panoptic_seg"]:
            evaluator_list.append(COCOEvaluator(dataset_name, output_dir=output_folder))
        if evaluator_type == "coco_panoptic_seg":
            evaluator_list.append(COCOPanopticEvaluator(dataset_name, output_folder))
        if evaluator_type == "cityscapes_instance":
            assert (
                torch.cuda.device_count() >= comm.get_rank()
            ), "CityscapesEvaluator currently do not work with multiple machines."
            return CityscapesInstanceEvaluator(dataset_name)
        if evaluator_type == "cityscapes_sem_seg":
            assert (
                torch.cuda.device_count() >= comm.get_rank()
            ), "CityscapesEvaluator currently do not work with multiple machines."
            return CityscapesSemSegEvaluator(dataset_name)
        elif evaluator_type == "pascal_voc":
            return PascalVOCDetectionEvaluator(dataset_name)
        elif evaluator_type == "lvis":
            return LVISEvaluator(dataset_name, output_dir=output_folder)
        if len(evaluator_list) == 0:
            raise NotImplementedError(
                "no Evaluator for the dataset {} with the type {}".format(
                    dataset_name, evaluator_type
                )
            )
        elif len(evaluator_list) == 1:
            return evaluator_list[0]
        return DatasetEvaluators(evaluator_list)

    @classmethod
    def test_with_TTA(cls, cfg, model):
        logger = logging.getLogger("detectron2.trainer")
        # In the end of training, run an evaluation with TTA
        # Only support some R-CNN models.
        logger.info("Running inference with test-time augmentation ...")
        model = GeneralizedRCNNWithTTA(cfg, model)
        evaluators = [
            cls.build_evaluator(
                cfg, name, output_folder=os.path.join(cfg.OUTPUT_DIR, "inference_TTA")
            )
            for name in cfg.DATASETS.TEST
        ]
        res = cls.test(cfg, model, evaluators)
        res = OrderedDict({k + "_TTA": v for k, v in res.items()})
        return res


def setup(args):
    """
    Create configs and perform basic setups.
    """
    cfg = get_cfg()
    cfg.merge_from_file(args.config_file)
    cfg.merge_from_list(args.opts)

    cfg.DATASETS.TRAIN = ("coco_my_train",) # 訓練資料集名稱,修改
    cfg.DATASETS.TEST = ("coco_my_val",) # 訓練資料集名稱,修改
    cfg.MODEL.RETINANET.NUM_CLASSES = 4 # 修改自己的類別數

    cfg.freeze()
    default_setup(cfg, args)
    return cfg


def main(args):
    cfg = setup(args)
    plain_register_dataset() #  # 修改


    if args.eval_only:
        model = Trainer.build_model(cfg)
        DetectionCheckpointer(model, save_dir=cfg.OUTPUT_DIR).resume_or_load(
            cfg.MODEL.WEIGHTS, resume=args.resume
        )
        res = Trainer.test(cfg, model)
        if cfg.TEST.AUG.ENABLED:
            res.update(Trainer.test_with_TTA(cfg, model))
        if comm.is_main_process():
            verify_results(cfg, res)
        return res

    """
    If you'd like to do anything fancier than the standard training logic,
    consider writing your own training loop (see plain_train_net.py) or
    subclassing the trainer.
    """
    trainer = Trainer(cfg)
    trainer.resume_or_load(resume=args.resume)
    if cfg.TEST.AUG.ENABLED:
        trainer.register_hooks(
            [hooks.EvalHook(0, lambda: trainer.test_with_TTA(cfg, trainer.model))]
        )
    return trainer.train()


if __name__ == "__main__":
    args = default_argument_parser().parse_args()
    print("Command Line Args:", args)
    launch(
        main,
        args.num_gpus,
        num_machines=args.num_machines,
        machine_rank=args.machine_rank,
        dist_url=args.dist_url,
        args=(args,),
    )

還有一些引數的修改也可以在def setup(args)中修改。

def setup(args):
    """
    Create configs and perform basic setups.
    """
    cfg = get_cfg()
    args.config_file = "../configs/COCO-Detection/mask_rcnn_R_50_FPN_3x.yaml"
    cfg.merge_from_file(args.config_file)   # 從config file 覆蓋配置
    cfg.merge_from_list(args.opts)          # 從CLI引數 覆蓋配置

    # 更改配置引數
    cfg.DATASETS.TRAIN = ("coco_my_train",) # 訓練資料集名稱
    cfg.DATASETS.TEST = ("coco_my_val",)
    cfg.DATALOADER.NUM_WORKERS = 4  # 單執行緒

    cfg.INPUT.CROP.ENABLED = True
    cfg.INPUT.MAX_SIZE_TRAIN = 640 # 訓練圖片輸入的最大尺寸
    cfg.INPUT.MAX_SIZE_TEST = 640 # 測試資料輸入的最大尺寸
    cfg.INPUT.MIN_SIZE_TRAIN = (512, 768) # 訓練圖片輸入的最小尺寸,可以設定為多尺度訓練
    cfg.INPUT.MIN_SIZE_TEST = 640
    #cfg.INPUT.MIN_SIZE_TRAIN_SAMPLING,其存在兩種配置,分別為 choice 與 range :
    # range 讓影象的短邊從 512-768隨機選擇
    #choice : 把輸入影象轉化為指定的,有限的幾種圖片大小進行訓練,即短邊只能為 512或者768
    cfg.INPUT.MIN_SIZE_TRAIN_SAMPLING = 'range'
#  本句一定要看下注釋!!!!!!!!
    cfg.MODEL.RETINANET.NUM_CLASSES = 81  # 類別數+1(因為有background,也就是你的 cate id 從 1 開始,如果您的資料集Json下標從 0 開始,這個改為您對應的類別就行,不用再加背景類!!!!!)
    #cfg.MODEL.WEIGHTS="/home/yourstorePath/.pth"
    cfg.MODEL.WEIGHTS = "/root/xxx/model_final_5bd44e.pkl"    # 預訓練模型權重
    cfg.SOLVER.IMS_PER_BATCH = 4  # batch_size=2; iters_in_one_epoch = dataset_imgs/batch_size

    # 根據訓練資料總數目以及batch_size,計算出每個epoch需要的迭代次數
    #9000為你的訓練資料的總數目,可自定義
    ITERS_IN_ONE_EPOCH = int(9000 / cfg.SOLVER.IMS_PER_BATCH)

    # 指定最大迭代次數
    cfg.SOLVER.MAX_ITER = (ITERS_IN_ONE_EPOCH * 12) - 1 # 12 epochs,
    # 初始學習率
    cfg.SOLVER.BASE_LR = 0.002
    # 優化器動能
    cfg.SOLVER.MOMENTUM = 0.9
    #權重衰減
    cfg.SOLVER.WEIGHT_DECAY = 0.0001
    cfg.SOLVER.WEIGHT_DECAY_NORM = 0.0
    # 學習率衰減倍數
    cfg.SOLVER.GAMMA = 0.1
    # 迭代到指定次數,學習率進行衰減
    cfg.SOLVER.STEPS = (7000,)
    # 在訓練之前,會做一個熱身運動,學習率慢慢增加初始學習率
    cfg.SOLVER.WARMUP_FACTOR = 1.0 / 1000
    # 熱身迭代次數
    cfg.SOLVER.WARMUP_ITERS = 1000

    cfg.SOLVER.WARMUP_METHOD = "linear"
    # 儲存模型檔案的命名資料減1
    cfg.SOLVER.CHECKPOINT_PERIOD = ITERS_IN_ONE_EPOCH - 1

    # 迭代到指定次數,進行一次評估
    cfg.TEST.EVAL_PERIOD = ITERS_IN_ONE_EPOCH
    #cfg.TEST.EVAL_PERIOD = 100

    #cfg.merge_from_file(args.config_file)
    #cfg.merge_from_list(args.opts)
    cfg.freeze()
    default_setup(cfg, args)
    return cfg

2.修改detectron2/detectron2/data/datasets/builtin.py(一般這個檔案不需要修改只是說明可以根據自己的路徑和命名來修改)

_PREDEFINED_SPLITS_COCO = {}

_PREDEFINED_SPLITS_COCO["coco"] = {
    "coco_2014_train": ("coco/train2014", "coco/annotations/instances_train2014.json"),
    "coco_2014_val": ("coco/val2014", "coco/annotations/instances_val2014.json"),
    "coco_2014_minival": ("coco/val2014", "coco/annotations/instances_minival2014.json"),
    "coco_2014_minival_100": ("coco/val2014", "coco/annotations/instances_minival2014_100.json"),
    "coco_2014_valminusminival": (
        "coco/val2014",
        "coco/annotations/instances_valminusminival2014.json",
    ),
    "coco_2017_train": ("coco/train2017", "coco/annotations/instances_train2017.json"),
    "coco_2017_val": ("coco/val2017", "coco/annotations/instances_val2017.json"),
    "coco_2017_test": ("coco/test2017", "coco/annotations/image_info_test2017.json"),
    "coco_2017_test-dev": ("coco/test2017", "coco/annotations/image_info_test-dev2017.json"),
    "coco_2017_val_100": ("coco/val2017", "coco/annotations/instances_val2017_100.json"),
    
	#我新註冊的資料集,注意路徑別寫錯了,以及註冊的名字如coco_raw_train別寫錯了   
    "coco_raw_train": ("coco_raw/train2017", "coco_raw/annotations/raw_train.json"),
    "coco_raw_val": ("coco_raw/val2017", "coco_raw/annotations/raw_val.json"),
    #"coco_raw_frcnn_train": ("coco_raw/train2017", "coco_raw/annotations/raw_train.json"),
    #"coco_raw_frcnn_val": ("coco_raw/val2017", "coco_raw/annotations/raw_val.json"),    
}

如果修改了這個檔案就要修改detectron2/configs/COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml中的

_BASE_: "../Base-RCNN-FPN.yaml" # 需要修改
MODEL:
  WEIGHTS: "detectron2://ImageNetPretrained/MSRA/R-50.pkl"
  MASK_ON: True
  RESNETS:
    DEPTH: 50
SOLVER:
  STEPS: (21000, 25000)
  MAX_ITER: 27000

修改detectron2/configs/Base-RCNN-FPN.yaml

MODEL:
  META_ARCHITECTURE: "GeneralizedRCNN"
  BACKBONE:
    NAME: "build_resnet_fpn_backbone"
  RESNETS:
    OUT_FEATURES: ["res2", "res3", "res4", "res5"]
  FPN:
    IN_FEATURES: ["res2", "res3", "res4", "res5"]
  ANCHOR_GENERATOR:
    SIZES: [[32], [64], [128], [256], [512]]  # One size for each in feature map
    ASPECT_RATIOS: [[0.5, 1.0, 2.0]]  # Three aspect ratios (same for all in feature maps)
  RPN:
    IN_FEATURES: ["p2", "p3", "p4", "p5", "p6"]
    PRE_NMS_TOPK_TRAIN: 2000  # Per FPN level
    PRE_NMS_TOPK_TEST: 1000  # Per FPN level
    # Detectron1 uses 2000 proposals per-batch,
    # (See "modeling/rpn/rpn_outputs.py" for details of this legacy issue)
    # which is approximately 1000 proposals per-image since the default batch size for FPN is 2.
    POST_NMS_TOPK_TRAIN: 1000
    POST_NMS_TOPK_TEST: 1000
  ROI_HEADS:
    NAME: "StandardROIHeads"
    IN_FEATURES: ["p2", "p3", "p4", "p5"]
  ROI_BOX_HEAD:
    NAME: "FastRCNNConvFCHead"
    NUM_FC: 2
    POOLER_RESOLUTION: 7
  ROI_MASK_HEAD:
    NAME: "MaskRCNNConvUpsampleHead"
    NUM_CONV: 4
    POOLER_RESOLUTION: 14
DATASETS:
  TRAIN: ("coco_raw_train",) # 修改與上面一致
  TEST: ("coco_raw_val",) # 修改與上面一致
SOLVER:
  IMS_PER_BATCH: 16
  BASE_LR: 0.02
  STEPS: (60000, 80000)
  MAX_ITER: 90000
INPUT:
  MIN_SIZE_TRAIN: (640, 672, 704, 736, 768, 800)
VERSION: 2

3.修改detectron2/detectron2/data/datasets/builtin_meta.py

COCO_CATEGORIES = [
    {"color": [220, 20, 60], "isthing": 1, "id": 1, "name": "class1"},
    {"color": [119, 11, 32], "isthing": 1, "id": 2, "name": "class2"},
    {"color": [0, 0, 142], "isthing": 1, "id": 3, "name": "class3"},
    {"color": [0, 0, 230], "isthing": 1, "id": 4, "name": "class4"},
    ...

根據自己的類別修改“name”:取值,或者在多餘81類後則可以在末尾新增,不然訓練出來的類別就顯示之前的類別了。
4.訓練

python tools/train_net.py \
--num-gpus 1 \
--config-file configs/COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml SOLVER.IMS_PER_BATCH 2 SOLVER.BASE_LR 0.0025

訓練的模型儲存在detectron2/output
5.測試

python tools/train_net.py \
--config-file configs/COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml \
--eval-only MODEL.WEIGHTS output/model_final.pth