facebook-pytorch maskrcnn-benchmark训练自己的数据集

  1. 数据准备
    gray标签图转coco数据集 https://www.cnblogs.com/yanghailin/p/11215003.html
    voc定位数据xml转coco数据集格式json https://www.cnblogs.com/yanghailin/p/11189871.html
    我把数据准备成coco格式的。在maskrcnn-benchmark根目录下面的datasets文件夹下,目录结构如下:
└── coco
    ├── annotations
    │   ├── instances_train2014.json
    │   └── instances_val2014.json
    ├── test2014
    │   └── cjh_993.png
    ├── train2014
    │   ├── cjh_968.png
    │   ├── cjh_969.png
    │   ├── cjh_976.png
    │   ├── cjh_977.png
    │   └── cjh_984.png
    └── val2014
        ├── cjh_985.png
        └── cjh_992.png

  1. 配置文件修改
    根目录下的configs文件夹里面有很多yaml文件,我选择的是e2e_mask_rcnn_R_50_FPN_1x.yaml,根据我的需求我更改如下:
MODEL:
  META_ARCHITECTURE: "GeneralizedRCNN"
  WEIGHT: "catalog://ImageNetPretrained/MSRA/R-50"
  BACKBONE:
    CONV_BODY: "R-50-FPN"
  RESNETS:
    BACKBONE_OUT_CHANNELS: 256
  RPN:
    USE_FPN: True
    ANCHOR_STRIDE: (4, 8, 16, 32, 64)
    PRE_NMS_TOP_N_TRAIN: 2000
    PRE_NMS_TOP_N_TEST: 1000
    POST_NMS_TOP_N_TEST: 1000
    FPN_POST_NMS_TOP_N_TEST: 1000
  ROI_HEADS:
    USE_FPN: True
  ROI_BOX_HEAD:
    POOLER_RESOLUTION: 7
    POOLER_SCALES: (0.25, 0.125, 0.0625, 0.03125)
    POOLER_SAMPLING_RATIO: 2
    FEATURE_EXTRACTOR: "FPN2MLPFeatureExtractor"
    PREDICTOR: "FPNPredictor"
    #修改成自己任务所需要检测的类别数+1  我是33类   33+1
    NUM_CLASSES: 34
  ROI_MASK_HEAD:
    POOLER_SCALES: (0.25, 0.125, 0.0625, 0.03125)
    FEATURE_EXTRACTOR: "MaskRCNNFPNFeatureExtractor"
    PREDICTOR: "MaskRCNNC4Predictor"
    POOLER_RESOLUTION: 14
    POOLER_SAMPLING_RATIO: 2
    RESOLUTION: 28
    SHARE_BOX_FEATURE_EXTRACTOR: False
  MASK_ON: True
DATASETS:
###我用的coco2014数据结构目录####################################
##此处的名字与maskrcnn-benchmark/maskrcnn_benchmark/config/paths_catalog.py 下的名字对应
##"coco_2014_train": {
##            "img_dir": "coco/train2014",
##          "ann_file": "coco/annotations/instances_train2014.json"
## 和coco数据结构一致的话就不需要修改
  TRAIN: ("coco_2014_train","coco_2014_val")
  TEST: ("coco_2014_test",)
DATALOADER:
  SIZE_DIVISIBILITY: 32
SOLVER:
  BASE_LR: 0.02
  WEIGHT_DECAY: 0.0001
  STEPS: (60000, 80000)
  MAX_ITER: 90000
  ###下面是我新增的##CHECKPOINT_PERIOD:500 训练每迭代500次保存一次模型 #####
## OUTPUT_DIR: 保存模型路径
  IMS_PER_BATCH: 3
  CHECKPOINT_PERIOD: 500
OUTPUT_DIR: "./my_weights/chejiahao/"
  1. maskrcnn-benchmark/maskrcnn_benchmark/config/defaults.py 改一处类别数改成自己的
    .#大概在214行 就改这一处就可以了 之前是81
_C.MODEL.ROI_BOX_HEAD.NUM_CLASSES = 34

4.训练
在根目录下运行:

python tools/train_net.py --config-file ./demo/my/e2e_mask_rcnn_R_50_FPN_1x.yaml SOLVER.IMS_PER_BATCH 2 SOLVER.BASE_LR 0.0025 SOLVER.MAX_ITER 720000 SOLVER.STEPS "(480000, 640000)" TEST.IMS_PER_BATCH 1 MODEL.RPN.FPN_POST_NMS_TOP_N_TRAIN 2000

就可以训练起来。训练显示的日志:

2019-07-19 17:11:18,959 maskrcnn_benchmark.trainer INFO: Start training
2019-07-19 17:11:30,359 maskrcnn_benchmark.trainer INFO: eta: 4 days, 17:59:34  iter: 20  loss: 2.7949 (3.2260)  loss_classifier: 0.9784 (1.3514)  loss_box_reg: 0.2458 (0.1996)  loss_mask: 0.7680 (0.9672)  loss_objectness: 0.6078 (0.6162)  loss_rpn_box_reg: 0.0833 (0.0916)  time: 0.5245 (0.5700)  data: 0.0089 (0.0395)  lr: 0.000897  max mem: 2562
2019-07-19 17:11:41,858 maskrcnn_benchmark.trainer INFO: eta: 4 days, 18:29:10  iter: 40  loss: 2.2825 (2.7907)  loss_classifier: 0.9786 (1.1715)  loss_box_reg: 0.3196 (0.2657)  loss_mask: 0.6615 (0.8187)  loss_objectness: 0.2881 (0.4529)  loss_rpn_box_reg: 0.0693 (0.0820)  time: 0.5792 (0.5725)  data: 0.0092 (0.0243)  lr: 0.000963  max mem: 2656

  1. 预测
    5.1. 修改之前配置的yaml文件e2e_mask_rcnn_R_50_FPN_1x.yaml的WEIGHT: 路径为自己训练好的模型路径/maskrcnn-benchmark/my_weights/chejiahao/model_0003500.pth
MODEL:
  META_ARCHITECTURE: "GeneralizedRCNN"
  WEIGHT: "maskrcnn-benchmark/my_weights/chejiahao/model_0003500.pth"
  BACKBONE:
    CONV_BODY: "R-50-FPN"

5.2 修改demo文件夹下面的predictor.py里面的类别标签列表为自己的

    # CATEGORIES = [
    #     "__background",
    #     "ZhiZaoNianYue",
    #     "FDJ_PaiL",
    #     "FDJ_Hao",
    #     "CJHao",
    # ]

在demo文件夹下面新建my_predictor.py

from maskrcnn_benchmark.config import cfg
from predictor import COCODemo
import cv2
import os

config_file = "maskrcnn-benchmark/demo/cjh/e2e_mask_rcnn_R_50_FPN_1x.yaml"

# update the config options with the config file
cfg.merge_from_file(config_file)
# manual override some options
cfg.merge_from_list(["MODEL.DEVICE", "cuda"])

coco_demo = COCODemo(
    cfg,
    min_image_size=800,
    confidence_threshold=0.7,
)

file_root = 'maskrcnn-benchmark/datasets/coco/test2014/'
file_list = os.listdir(file_root)
save_out = "maskrcnn-benchmark/demo/chejiahao/out/"
for img_name in file_list:
    img_path = file_root + img_name
    image = cv2.imread(img_path)
    predictions = coco_demo.run_on_opencv_image(image)
    save_path = save_out + img_name
    cv2.imwrite(save_path,predictions)

    cv2.namedWindow("img",cv2.WINDOW_NORMAL)
    cv2.imshow("img",predictions)
    cv2.waitKey(1)

在demo文件夹下运行:

python my_predictor.py

enjoy!

posted @ 2019-07-19 17:19  无左无右  阅读(2021)  评论(0编辑  收藏  举报