美文网首页深度学习-推荐系统-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训练自己的数据集

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

  • TX2 +caffe +ssd/ubuntu16.04+caff

    1.SSD安装及训练自己的数据集 - CSDN博客 2.caffe-SSD训练自己的数据集 - CSDN博客 3....

  • CS231N学习记录

    数据集:训练集+验证集+测试集 交叉验证:当训练数据太小时,为了更好地利用数据,那么将训练数据集划分成n份,其中n...

  • sklearn学习 — 数据集

    sklearn数据集 1. 数据集的划分 训练集 : (占数据集比重高) 用于训练,构建模型 测试集 : 在模型...

  • 神经网络优化1

    数据划分 数据集分类 通常会将数据集分层三类: 训练集(Training Sets):采用训练集进行训练时,通过改...

  • 2.封装kNN算法之数据分割

    训练数据集与测试数据集 当我们拿到一组数据之后,通常我们需要把数据分割成两部分,即训练数据集和测试数据集。训练数据...

  • 洛杉矶房价预测

    制作训练集、评测集 交叉验证 数据有限,发挥数据本来的效率 数据的训练集合评测集的矛盾a. 如果用更多的数据去训练...

  • 基于Keras实现Kaggle2013--Dogs vs. Ca

    【下载数据集】 下载链接--百度网盘关于猫的部分数据集示例 【整理数据集】 将训练数据集分割成训练集、验证集、测试...

  • Caffe学习记录02

    用AlexNet训练自己的数据集 参考资料 本篇重点在于讲述如何讲自己的数据集转化为Caffe能用的数据集,至于具...

  • 高数没用?看看机器学习的公式,验证下微积分都还给老师没

    KNN分类 K-近邻算法,把数据集分为训练数据集和测试数据集,用训练数据集导入KNN,对于一个数据向量,计算此数据...

网友评论

    本文标题:SimpleDet训练自己的数据集

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