CornerNer_Lite训练自己的数据集

0 前言

在UBUNTU16.04使用CornerNet_Lite进行目标识别并嵌入到ROS中一文中，介绍了如何使用CornerNet_Lite的两种方法检测目标和嵌入到ROS中使用。但是使用的是官方训练好的COCO数据集，有80个类，但实际应用中，可能到某个应用场景只需要几个类而已，多余的类返回会干扰造成误识别的情况。因此有必要训练出自己需要的数据集文件用于检测。

1 环境搭建

上文UBUNTU16.04使用CornerNet_Lite进行目标识别并嵌入到ROS中中提到了单纯应用所需要的环境，编译了Corner PoolingLayers和NMS，这里训练数据集的话还需要安装MS COCO API。这里再把完成过程叙述一遍。

1）github上下载源码

CornerNet_ite github网址

git clone https://github.com/princeton-vl/CornerNet-Lite.git

2）环境配置

(1）创建虚拟环境

cd <CornerNet-Lite dir>
conda create --name CornerNet_Lite --file conda_packagelist.txt --channel pytorch

创建CornerNet_Lite环境的过程需要下载一个1.0版本的pytorch有600多M,一次不一定能成功，失败了重新再运行下。

（2）激活环境

conda activate CornerNet_Lite

此时终端最前面个会多一个（CornerNet_LIte）的标志，表示已经进入到该虚拟空间。

退出虚拟空间：conda deactivate

（3）编译Corner PoolingLayers

cd <CornerNet-Lite dir>/core/models/py_utils/_cpools/
python setup.py install --user

(4)编译NMS

cd <CornerNet-Lite dir>/core/external
make

(5)安装MS COCO API

mkdir -p <CornerNet-Lite dir>/data
cd <CornerNet-Lite dir>/data
git clone git@github.com:cocodataset/cocoapi.git coco
cd <CornerNet-Lite dir>/data/coco/PythonAPI
make install

如果url ：git@github.com:cocodataset/cocoapi.git有问题，也可以从下述网址下载cocoAPI：https://github.com/cocodataset/cocoapi.git

git clone https://github.com/cocodataset/cocoapi.git coco

2 准备自己的数据集

之前都这里有做好的几个voc数据集，需要将voc转换为coco。

CornerNet-Lite默认代码通过cocoAPI加载coco json文件的形式来训练的，如果想训练voc格式的数据，有两种办法：一、修改代码加载标注文件形式，改为加载xml的格式。二、将xml转化为coco json格式的标注文件。这里将采用第二种方法。

首先注意图片名称和标注文件名称要统一，且coco中图片明明是按照image id的顺序排列的，注意名称中只能包含数字，这里建议重新统一命名，可采用如下代码(coco_rename.py)：

# coding utf-8
import os, shutil
import random
import numpy as np
 
def renameData(xmlDir, imgDir):
    xmlFiles = os.listdir(xmlDir)
    total = len(xmlFiles)
    cur = 0
    for xml in xmlFiles:
        cur += 1
        if cur % 500 == 1:
            print("Total/cur:", total, "/", cur)
        imgPath = imgDir + xml[:-4] + ".jpg"
        
        outName = ("%08d" % (cur))
        
        outXMLPath = ("%s/%s.xml" % (xmlDir,outName))
        print("xml transformed successed:",cur)
        outImgPath = ("%s/%s.jpg" % (imgDir,outName))
        print("img transformed successed:",cur)
        
        os.rename(xmlDir+xml,outXMLPath)
        os.rename(imgPath,outImgPath)
        
    
    print("picker number:",cur)
    
 
if __name__ == '__main__':
    #下面的图片路径和标注路径注意更改为自己的实际路径
    xmlDir = "/home/zbr/zbr_work/CornerNet-Lite/data/Annotations/"    
    imgDir = "/home/zbr/zbr_work/CornerNet-Lite/data/JPEGImages/"
    
    print(xmlDir)
    print(imgDir)
    
    renameData(xmlDir, imgDir)

命名数字保留了8位数，不足前面补0.

然后需要按照自己的需求将标注数据.xml文件分成train,val,test。对于voc数据集而言就是annotations文件夹所有的内容，代码如下（coco_splitDate.py），

# -*- coding:utf-8 -*-
 
import random
from random import randint
import os
import numpy as np
import shutil
 
def clc(dir):
    ls = os.listdir(dir)
    for i in ls:
        c_path = os.path.join(dir, i)
        os.remove(c_path)
 
if __name__ == '__main__':
    folder_list= ["train","val","test"]
    #注意更改base_dir为本地实际图像和标注文件路径
    base_dir = "/home/zbr/zbr_work/CornerNet-Lite/data/" 
    for i in range(3):
        folderName = folder_list[i]
        xml_dir = base_dir + folderName + "/Annotations/"
        if not os.path.exists(xml_dir):
        	os.makedirs(xml_dir)
 
    train_dir = base_dir + folder_list[0] + "/Annotations/"
    val_dir = base_dir + folder_list[1] + "/Annotations/"
    test_dir = base_dir + folder_list[2] + "/Annotations/"
    clc(train_dir)
    clc(test_dir)
    clc(val_dir)
 
    total = np.arange(1,121+1)  #total sample 121
    valList = random.sample(range(1, 121+1), 24) #val random select 20% of total
    print(sorted(valList))
 
    # move 24 val pics to splited folder
    for i in range(0,24):
    	xmlfile = ("%08d" % (valList[i]))
    	file_dir = base_dir + "Annotations"+'/'+ xmlfile +'.xml'
    	shutil.copy(file_dir,val_dir)
 
    # take away val ids from total ids
    total_val = [] # total ids minus val ids, ids left stay here
    for i in range(0,121):
    	if total[i] not in valList:
    		total_val.append(total[i])
 
    testList_temp = random.sample(range(0, 121-24), 24) #test random select 20% of total
    print(sorted(testList_temp))
    testList = []
    for k in range(