美文网首页深度学习-推荐系统-CV-NLP
SimpleDet训练自己的数据集

SimpleDet训练自己的数据集

作者: cshun | 来源:发表于2019-08-07 15:37 被阅读2次

    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训练自己的数据集

        本文链接:https://www.haomeiwen.com/subject/fxbodctx.html