SimpleDet训练自己的数据集
2019-08-07 本文已影响2人
cshun
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
