SimpleDet训练自己的数据集

SimpleDet是一套简单通用的目标检测与物体识别的框架。整套框架基于MXNet的原生API完成。

github:https://github.com/TuSimple/simpledet](https://github.com/TuSimple/simpledet

框架作者分享(知乎):SimpleDet: 一套简单通用的目标检测与物体识别框架

使用docker配置环境(这个需要服务器上安装好nvidia-docker):

pre-built docker images for both cuda9.0 and cuda10.0.
Maxwell, Pascal, Volta and Turing GPUs are supported.
For nvidia-driver >= 410.48, cuda10 image is recommended.
For nvidia-driver >= 384.81, cuda9 image is recommended.
Aliyun beijing mirror is provided for users pulling from China.

nvidia-docker run -it -v $HOST-SIMPLEDET-DIR:$CONTAINER-WORKDIR rogerchen/simpledet:cuda9 zsh
nvidia-docker run -it -v $HOST-SIMPLEDET-DIR:$CONTAINER-WORKDIR rogerchen/simpledet:cuda10 zsh
nvidia-docker run -it -v $HOST-SIMPLEDET-DIR:$CONTAINER-WORKDIR registry.cn-beijing.aliyuncs.com/rogerchen/simpledet:cuda9 zsh
nvidia-docker run -it -v $HOST-SIMPLEDET-DIR:$CONTAINER-WORKDIR registry.cn-beijing.aliyuncs.com/rogerchen/simpledet:cuda10 zsh

过程中CONTAINER-WORKDIR可能需要设置绝对路径,来完成挂载。否则会报错。

其他环境配置方法参考可https://github.com/TuSimple/simpledet/blob/master/doc/INSTALL.md

配置其他依赖:mxnext和others

1.setup mxnext, a wrapper of mxnet symbolic API

cd $SIMPLEDET_DIR
git clone https://github.com/RogerChern/mxnext

2.run make in simpledet directory to install cython extensions

make 

数据格式:

[
    {
        "gt_class": (nBox, ),
        "gt_bbox": (nBox, 4),
        "flipped": bool,
        "h": int,
        "w": int,
        "image_url": str,
        "im_id": int,

        # this fields are generated on the fly during test
        "rec_id": int,
        "resize_h": int,
        "resize_w": int,
        ...
    },
    ...
]

在处理数据过程中,可以使用coco格式数据,组织方式如下:

data/
    coco/
        annotations/
            instances_train2014.json
            instances_valminusminival2014.json
            instances_minival2014.json
            image_info_test-dev2017.json
        images/
            train2014
            val2014
            test2017

之后运行下述命令生成索引

python3 utils/generate_roidb.py --dataset coco --dataset-split train2014
python3 utils/generate_roidb.py --dataset coco --dataset-split valminusminival2014
python3 utils/generate_roidb.py --dataset coco --dataset-split minival2014
python3 utils/generate_roidb.py --dataset coco --dataset-split test-dev2017

在generate_roidb.py文件中可以指定data/coco/annotations/*.json文件对应的data/coco/image/下文件夹。

配置好数据,就可以愉快的训练了。

训练

# train
python3 detection_train.py --config config/detection_config.py

# test
python3 detection_test.py --config config/detection_config.py

在这个过程中,挑选好合适的config文件,并对config/detection_config.py文件进行编辑。

遇到的一个问题和解决方案:

在训练集中有部分单张图上目标数量超过了100,但是config中max_num_gt=100,导致数据读取过程出现错误,表现形式是在训练过程中突然卡顿并不显示错误。
解决方案:将max_num_gt=300。

加载权重

在MODEL_ZOO.md中可以寻找合适的权重信息

代码结构如下:

detection_train.py
detection_test.py
config/
    detection_config.py
core/
    detection_input.py
    detection_metric.py
    detection_module.py
models/
    FPN/
    tridentnet/
    maskrcnn/
    cascade_rcnn/
    retinanet/
mxnext/
symbol/
    builder.py

训练后保存权重和log文件等在experiment文件夹下:
One experiment is a directory in experiments folder with the same name as the config file.
E.g. r50_fixbn_1x.py is the name of a config file

config/
    r50_fixbn_1x.py
experiments/
    r50_fixbn_1x/
        checkpoint.params
        log.txt(训练日志)
        coco_minival2014_result.json(运行detection_test.py后的结果文件。)

单张图像测试/可视化

在simpledet目录下新建detect_image.py文件。代码如下。
运行命令(image_path为图像路径):

python3 detect_image.py image_path
import cv2
import os
import argparse
import importlib
import mxnet as mx
import numpy as np

from core.detection_module import DetModule
from utils.load_model import load_checkpoint

CATEGORIES = [
    "__background",
    "airplane",
    "helicopter"
]

def parse_args():
    parser = argparse.ArgumentParser(description='Test Detection')
    # general
    parser.add_argument('img', help='the image path', type=str)
    parser.add_argument('--config', help='config file path', type=str, default='config/tridentnet_r50v1c4_c5_1x.py')
    parser.add_argument('--batch_size', help='', type=int, default=1)
    parser.add_argument('--gpu', help='the gpu id for inferencing', type=int, default=0)
    parser.add_argument('--thresh', help='the threshold for filtering boxes', type=float, default=0.7)
    args = parser.parse_args()

    return args


class predictor(object):
    def __init__(self, config, batch_size, gpu_id, thresh):
        self.config = config
        self.batch_size = batch_size
        self.thresh = thresh

        # Parse the parameter file of model
        pGen, pKv, pRpn, pRoi, pBbox, pDataset, pModel, pOpt, pTest, \
        transform, data_name, label_name, metric_list = config.get_config(is_train=False)

        self.data_name = data_name
        self.label_name = label_name
        self.p_long, self.p_short = transform[2].p.long, transform[2].p.short

        # Define NMS type
        if callable(pTest.nms.type):
            self.do_nms = pTest.nms.type(pTest.nms.thr)
        else:
            from operator_py.nms import py_nms_wrapper

            self.do_nms = py_nms_wrapper(pTest.nms.thr)

        sym = pModel.test_symbol
        sym.save(pTest.model.prefix + "_test.json")

        ctx = mx.gpu(gpu_id)
        data_shape = [
            ('data', (batch_size, 3, 800, 1200)),
            ("im_info", (1, 3)),
            ("im_id", (1,)),
            ("rec_id", (1,)),
        ]

        # Load network
        arg_params, aux_params = load_checkpoint(pTest.model.prefix, pTest.model.epoch)
        from utils.graph_optimize import merge_bn
        sym, arg_params, aux_params = merge_bn(sym, arg_params, aux_params)
        self.mod = DetModule(sym, data_names=data_name, context=ctx)
        self.mod.bind(data_shapes=data_shape, for_training=False)
        self.mod.set_params(arg_params, aux_params, allow_extra=False)

    def preprocess_image(self, input_img):
        image = input_img[:, :, ::-1]  # BGR -> RGB

        short = min(image.shape[:2])
        long = max(image.shape[:2])
        scale = min(self.p_short / short, self.p_long / long)

        h, w = image.shape[:2]
        im_info = (round(h * scale), round(w * scale), scale)

        image = cv2.resize(image, None, None, scale, scale, interpolation=cv2.INTER_LINEAR)
        image = image.transpose((2, 0, 1))  # HWC -> CHW

        return image, im_info

    def run_image(self, img_path):
        image = cv2.imread(img_path, cv2.IMREAD_COLOR)
        image, im_info = self.preprocess_image(image)
        input_data = {'data': [image],
                      'im_info': [im_info],
                      'im_id': [0],
                      'rec_id': [0],
                      }

        data = [mx.nd.array(input_data[name]) for name in self.data_name]
        label = []
        provide_data = [(k, v.shape) for k, v in zip(self.data_name, data)]
        provide_label = [(k, v.shape) for k, v in zip(self.label_name, label)]

        data_batch = mx.io.DataBatch(data=data,
                                     label=label,
                                     provide_data=provide_data,
                                     provide_label=provide_label)

        self.mod.forward(data_batch, is_train=False)
        out = [x.asnumpy() for x in self.mod.get_outputs()]

        cls_score = out[3]
        bboxes = out[4]

        result = {}
        for cid in range(cls_score.shape[1]):
            if cid == 0:  # Ignore the background
                continue
            score = cls_score[:, cid]
            if bboxes.shape[1] != 4:
                cls_box = bboxes[:, cid * 4:(cid + 1) * 4]
            else:
                cls_box = bboxes
            valid_inds = np.where(score >= self.thresh)[0]
            box = cls_box[valid_inds]
            score = score[valid_inds]
            det = np.concatenate((box, score.reshape(-1, 1)), axis=1).astype(np.float32)
            det = self.do_nms(det)
            if len(det) > 0:
                det[:, :4] = det[:, :4] / im_info[2]  # Restore to the original size
                result[CATEGORIES[cid]] = det

        return result


if __name__ == "__main__":
    os.environ["MXNET_CUDNN_AUTOTUNE_DEFAULT"] = "0"

    args = parse_args()
    img_path = args.img
    config = importlib.import_module(args.config.replace('.py', '').replace('/', '.'))
    batch_size = args.batch_size
    gpu_id = args.gpu
    thresh = args.thresh

    save_dir = 'out'
    if not os.path.exists(save_dir):
        os.makedirs(save_dir)

    coco_predictor = predictor(config, batch_size, gpu_id, thresh)
    result = coco_predictor.run_image(img_path)

    draw_img = cv2.imread(img_path)
    for k, v in result.items():
        print('%s, num:%d' % (k, v.shape[0]))
        for box in v:
            score = box[4]
            box = box.astype(int)
            x1, y1, x2, y2 = box[:4]

            cv2.putText(draw_img, '%s:%.2f' % (k, score), (x1, y1), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255))
            cv2.rectangle(draw_img, (x1, y1), (x2, y2), (255, 0, 0), 2)

    save_name = os.path.basename(img_path)
    cv2.imwrite(os.path.join(save_dir, 'result_%s' % save_name), draw_img)

参考:https://github.com/TuSimple/simpledet
https://blog.csdn.net/f16011/article/details/88785792

你可能感兴趣的:(SimpleDet训练自己的数据集)