【目标检测】使用Retinanet训练自己的数据集——Keras版本

---

Retinanet 较之 SSD, YOLO系one stage目标检测网络，在保证检测速度的基础上，很大的提高了在样本类别分布不平衡情况下的检测精度，这得益于He Kaiming等人所提出的Focal Loss
Paper可见

Focal Loss for Dense Object Detection

由于在工作或一些特定场景中，我们的数据很大可能不会像COCO, Pascal VOC这样的实验数据集有很多分布相对均匀的类，而基本是集中在某几种类别上。这时使用Retinanet 就比较合适了。
本文主要介绍如何使用Keras版本的Retinanet训练自己的数据集，代码及参考来自：

fizyr/keras-retinanet

---

我的开发环境 Ubuntu 16.04 + python 3.6
Windows下亦可运行，但需要Microsoft Visual C++ 14.0支持
下载：Microsoft Visual C++ Build Tools

---

1 准备自己的数据集

根据自己系统下载labelImg工具，对自己的图片数据打上标签。数据集可自行爬取相关类目，或这里推荐一位同事自己制作数据集的博客，有非常详细的目标检测数据准备步骤及数据下载

目标检测第2步-数据准备(简书作者：潇洒坤)

2 搭建Retinanet环境

代码库下载地址https://github.com/fizyr/keras-retinanet，或git命令：

git clone https://github.com/fizyr/keras-retinanet.git

1. 获得代码库后进入keras-retinanet文件夹
2. 确认有未安装numpy

pip install numpy --user

3. 在这个文件夹内运行下面代码来安装keras-retinanet库，确认你已经根据自己的系统需求安装了tensorflow

pip install . --user 

或者你可以选择不安装库，仅在这个文件内运行相关训练测试代码，但是要先运行下述命令来编译Cython代码*(尽量不选择这种方式，仅能在此文件夹内运行相关代码，非常麻烦）*

python setup.py build_ext --inplace

3 生成CSV文件定义自己的数据集

训练自己的数据集需要至少两个CSV文件，一个文件包含标注数据，另一个则包含各个类别名及其对应的ID序号映射。

• 数据标注文件格式

  路径,xmin,ymin,xmax,ymax,类别名
  path/to/image.jpg,x1,y1,x2,y2,class_name
  

  注：数据标注的CSV文件需一行仅包含一条标注信息，这和yolo标注数据集形成的.txt有所不同，即若一幅图片包含了多个标注的bounding boxes ，一行也仅显示一条标注，共分几行来进行标识，若没有则路径后仅有‘，’，示例如下：

  /data/imgs/img_001.jpg,837,346,981,456,cow
  /data/imgs/img_002.jpg,215,312,279,391,cat
  /data/imgs/img_002.jpg,22,5,89,84,bird
  /data/imgs/img_003.jpg,,,,,
  

• 类别映射文件格式

  类别，序号（从0开始）
  class_name,id
  例：
  cow,0
  cat,1
  bird,2
  

写了个代码根据xml生成csv文件，下载：xml2csv.py

"""
运行方式：命令行 python xml2csv.py -i indir(图片及标注的母目录)
      注：必须参数: -i 指定包含有图片及标注的母文件夹，图片及标注可不在同一子目录里，但名称必须一一对应
                     (图片格式默认.jpg,若为其他格式可见代码中注释自行修改)
          可选参数: -p 交叉验证集拆分比，默认0.05
                   -t 生成训练集CSV文件名称，默认train.csv
                   -v 生成交叉验证集CSV文件名称，默认val.csv
                   -c 生成类别映射CSV文件名称，默认class.csv
"""

import os
import xml.etree.ElementTree as ET
import random
import math
import argparse


def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument('-i', '--indir', type=str)
    parser.add_argument('-p', '--percent', type=float, default=0.05)
    parser.add_argument('-t', '--train', type=str, default='train.csv')
    parser.add_argument('-v', '--val', type=str, default='val.csv')
    parser.add_argument('-c', '--classes', type=str, default='class.csv')
    args = parser.parse_args()
    return args

#获取特定后缀名的文件列表
def get_file_index(indir, postfix):
    file_list = []
    for root, dirs, files in os.walk(indir):
        for name in files:
            if postfix in name:
                file_list.append(os.path.join(root, name))
    return file_list

#写入标注信息
def convert_annotation(csv, address_list):
    cls_list = []
    with open(csv, 'w') as f:
        for i, address in enumerate(address_list):
            in_file = open(address, encoding='utf8')
            strXml =in_file.read()
            in_file.close()
            root=ET.XML(strXml)
            for obj in root.iter('object'):
                cls = obj.find('name').text
                cls_list.append(cls)
                xmlbox = obj.find('bndbox')
                b = (int(xmlbox.find('xmin').text), int(xmlbox.find('ymin').text), 
                     int(xmlbox.find('xmax').text), int(xmlbox.find('ymax').text))
                f.write(file_dict[address_list[i]])
                f.write( "," + ",".join([str(a) for a in b]) + ',' + cls)
                f.write('\n')
    return cls_list


if __name__ == "__main__":
    args = parse_args()
    file_address = args.indir
    test_percent = args.percent
    train_csv = args.train
    test_csv = args.val
    class_csv = args.classes
    Annotations = get_file_index(file_address, '.xml')
    Annotations.sort()
    JPEGfiles = get_file_index(file_address, '.jpg') #可根据自己数据集图片后缀名修改
    JPEGfiles.sort()
    assert len(Annotations) == len(JPEGfiles) #若XML文件和图片文件名不能一一对应即报错
    file_dict = dict(zip(Annotations, JPEGfiles))
    num = len(Annotations)
    test = random.sample(k=math.ceil(num*test_percent), population=Annotations)
    train = list(set(Annotations) - set(test))

    cls_list1 = convert_annotation(train_csv, train)
    cls_list2 = convert_annotation(test_csv, test)
    cls_unique = list(set(cls_list1+cls_list2))

    with open(class_csv, 'w') as f:
        for i, cls in enumerate(cls_unique):
            f.write(cls + ',' + str(i) + '\n')

4 开始训练

基础命令：

#若你已经安装了Retinanet库，在生成的csv文件目录下即可直接运行:
retinanet-train csv 标注csv文件路径(./train.csv)  类别映射csv文件路径(./class.csv) 
#若直接在keras-retainanet文件夹中运行
keras_retinanet/bin/train.py csv 标注csv文件路径  类别映射csv文件路径

更多的参数设定可参见keras_retinanet/bin/train.py
与训练自己数据集有关的参数如下：

"""
csv 后第一个参数接标注csv文件路径 ， 第二个接类别映射csv文件路径， 第三个参数可选择添加交叉验证集
示例：retinanet-train csv ./train.csv ./class.csv --val-annotations ./val.csv

一般还需指定 --epochs 训练轮数 默认值50
            --batch-size 一批训练多少个 默认值1
            --steps 一轮训练多少步 默认10000 需按照自己数据集size大小计算 steps = size / batch-size

若有12000个样本， batch_size=4，训练20轮， 则命令如下：
retinanet-train --batch-size 4 --steps 3000 --epochs 20 csv ./train.csv ./class.csv --val-annotations ./val.csv

至于是否加载权重训练，backbone选择（默认Resnet50，可选参见keras_retinanet/models），学习率大小等按照自己需要进行指定。
若要多GPU训练，经测试指定multi-gpu数量提示该功能处于测试阶段，若一定要使用，需指定multi-gpu-force, 后果自行承担~

"""

    csv_parser = subparsers.add_parser('csv')
    csv_parser.add_argument('annotations', help='Path to CSV file containing annotations for training.')
    csv_parser.add_argument('classes', help='Path to a CSV file containing class label mapping.')
    csv_parser.add_argument('--val-annotations', help='Path to CSV file containing annotations for validation (optional).')

    group = parser.add_mutually_exclusive_group()
    group.add_argument('--snapshot',          help='Resume training from a snapshot.')
    group.add_argument('--imagenet-weights',  help='Initialize the model with pretrained imagenet weights. This is the default behaviour.', action='store_const', const=True, default=True)
    group.add_argument('--weights',           help='Initialize the model with weights from a file.')
    group.add_argument('--no-weights',        help='Don\'t initialize the model with any weights.', dest='imagenet_weights', action='store_const', const=False)

    parser.add_argument('--backbone',         help='Backbone model used by retinanet.', default='resnet50', type=str)
    parser.add_argument('--batch-size',       help='Size of the batches.', default=1, type=int)
    parser.add_argument('--gpu',              help='Id of the GPU to use (as reported by nvidia-smi).')
    parser.add_argument('--multi-gpu',        help='Number of GPUs to use for parallel processing.', type=int, default=0)
    parser.add_argument('--multi-gpu-force',  help='Extra flag needed to enable (experimental) multi-gpu support.', action='store_true')
    parser.add_argument('--epochs',           help='Number of epochs to train.', type=int, default=50)
    parser.add_argument('--steps',            help='Number of steps per epoch.', type=int, default=10000)
    parser.add_argument('--lr',               help='Learning rate.', type=float, default=1e-5)
    parser.add_argument('--snapshot-path',    help='Path to store snapshots of models during training (defaults to \'./snapshots\')', default='./snapshots')
    parser.add_argument('--tensorboard-dir',  help='Log directory for Tensorboard output', default='./logs')
    parser.add_argument('--no-snapshots',     help='Disable saving snapshots.', dest='snapshots', action='store_false')
    parser.add_argument('--no-evaluation',    help='Disable per epoch evaluation.', dest='evaluation', action='store_false')
    parser.add_argument('--freeze-backbone',  help='Freeze training of backbone layers.', action='store_true')
    parser.add_argument('--random-transform', help='Randomly transform image and annotations.', action='store_true')
    parser.add_argument('--image-min-side',   help='Rescale the image so the smallest side is min_side.', type=int, default=800)
    parser.add_argument('--image-max-side',   help='Rescale the image if the largest side is larger than max_side.', type=int, default=1333)
    parser.add_argument('--config',           help='Path to a configuration parameters .ini file.')
    parser.add_argument('--weighted-average', help='Compute the mAP using the weighted average of precisions among classes.', action='store_true')
    parser.add_argument('--compute-val-loss', help='Compute validation loss during training', dest='compute_val_loss', action='store_false')

    # Fit generator arguments
    parser.add_argument('--workers', help='Number of multiprocessing workers. To disable multiprocessing, set workers to 0', type=int, default=1)
    parser.add_argument('--max-queue-size', help='Queue length for multiprocessing workers in fit generator.', type=int, default=10)

训练时若没有Resnet50的Backbone会先进行下载, 待下载结束模型构架完成即开始训练，默认训练好一轮保存一次模型，训练好模型文件姜保存在snapshot文件夹中，按照训练的步数进行命名，如 resnet_csv_15.h5, 注意：这里训练好的模型不能直接用于目标检测，需要进行转换。

5 将训练好模型转化为推断模型

命令行运行：

# 若你已经安装了Retinanet库
retinanet-convert-model 训练出的模型地址 转化后的推断模型地址
retinanet-convert-model /path/to/training/model.h5 /path/to/save/inference/model.h5
# 若直接在文件夹中运行
keras_retinanet/bin/convert_model.py /path/to/training/model.h5 /path/to/save/inference/model.h5

6 模型测试

示例代码在keras-retinanet中examples中，亦可见keras-retinanet/examples/ResNet50RetinaNet.ipynb

一般使用下述代码即可预测：

from keras_retinanet.models import load_model

#需要填入转化后的推断模型地址
model = load_model('/path/to/model.h5', backbone_name='resnet50')

boxes, scores, labels = model.predict_on_batch(inputs)

预测图片：

%matplotlib inline
%load_ext autoreload
%autoreload 2

import keras

from ..keras_retinanet import models
from ..keras_retinanet.utils.image import read_image_bgr, preprocess_image, resize_image
from ..keras_retinanet.utils.visualization import draw_box, draw_caption
from ..keras_retinanet.utils.colors import label_color

import matplotlib.pyplot as plt
import cv2
import os
import numpy as np
import time

import tensorflow as tf

def get_session():
    config = tf.ConfigProto()
    config.gpu_options.allow_growth = True
    return tf.Session(config=config)

# 选择要使用的GPU flag
#os.environ["CUDA_VISIBLE_DEVICES"] = "1"

# 设置tensorflow session 为Keras 后端
keras.backend.tensorflow_backend.set_session(get_session())
#转化后的推断模型地址
model_path = os.path.join('..', 'snapshots', 'predict.h5')
#加载模型
model = models.load_model(model_path, backbone_name='resnet50')
#建立ID与类别映射字典
labels_to_names = {0: 'class_1', 1: 'class_2',...}
#加载需要检测的图片
image_path = ' '
image = read_image_bgr(image_path)

# copy到另一个对象并转为RGB文件
draw = image.copy()
draw = cv2.cvtColor(draw, cv2.COLOR_BGR2RGB)

# 图像预处理
image = preprocess_image(image)
image, scale = resize_image(image)

# 模型预测
start = time.time()
boxes, scores, labels = model.predict_on_batch(np.expand_dims(image, axis=0))
print("processing time: ", time.time() - start)

# 矫正比例
boxes /= scale

# 目标检测可视化展示
for box, score, label in zip(boxes[0], scores[0], labels[0]):
    # 设置预测得分最低阈值
    if score < 0.5:
        break
        
    color = label_color(label)
    
    b = box.astype(int)
    draw_box(draw, b, color=color)
    
    caption = "{} {:.3f}".format(labels_to_names[label], score)
    draw_caption(draw, b, caption)
    
#图片展示
plt.figure(figsize=(15, 15))
plt.axis('off')
plt.imshow(draw)
plt.show()

初发布于我的简书