[转载]使用ImageNet在faster-rcnn上训练自己的分类网络

原文链接 www.cnblogs.com/zhiyishou/p/5651321.html

具体代码见https://github.com/zhiyishou/py-faster-rcnn

[转载]使用ImageNet在faster-rcnn上训练自己的分类网络_第1张图片

这是我对cup, glasses训练的识别

faster-rcnn在fast-rcnn的基础上加了rpn来将整个训练都置于GPU内,以用来提高效率,这里我们将使用ImageNet的数据集来在faster-rcnn上来训练自己的分类器。从ImageNet上可下载到很多类别的Image与bounding box annotation来进行训练(每一个类别下的annotation都少于等于image的个数,所以我们从annotation来建立索引)。

在lib/dataset/factory.py中提供了coco与voc的数据集获取方法,而我们要做的就是在这里加上我们自己的ImageNet获取方法,我们先来建立ImageNet数据获取主文件。coco与pascal_voc的获取都是继承于父类imdb,所以我们可根据pascal_voc的获取方法来做模板修改完成我们的ImageNet类。

创建ImageNet类

由于在faster-rcnn里使用rpn来代替了selective_search,所以我们可以在使用时直接略过有关selective_search的方法,根据pascal_voc类做模板,我们需要留下的方法有:

__init__//初始化image_path_at//根据数据集列表的index来取图片绝对地址image_path_from_index//配合上面_load_image_set_index//获取数据集列表_gt_roidb//获取ground-truth数据rpn_roidb//获取region proposal数据_load_rpn_roidb//根据gt_roidb生成rpn_roidb数据并合成_load_psacal_annotation//加载annotation文件并对bounding box进行数据整理

__init__:

def __init__(self, image_set):        imdb.__init__(self,'imagenet')self._image_set = image_setself._data_path = os.path.join(cfg.DATA_DIR,"imagenet")#类别与对应的wnid,可以修改成自己要训练的类别self._class_wnids = {'cup':'n03147509','glasses':'n04272054'}#类别,修改类别时同时要修改这里self._classes = ('__background__',self._class_wnids['cup'],self._class_wnids['glasses'])self._class_to_ind = dict(zip(self.classes, xrange(self.num_classes)))#bounding box annotation 文件的目录self._xml_path = os.path.join(self._data_path,"Annotations")self._image_ext ='.JPEG'#我们使用xml文件名来做数据集的索引# the xml file name and each one corresponding to image file nameself._image_index =self._load_xml_filenames()self._salt = str(uuid.uuid4())self._comp_id ='comp4'self.config = {'cleanup':True,'use_salt':True,'use_diff':False,'matlab_eval':False,'rpn_file': None,'min_size':2}        assert os.path.exists(self._data_path), \'Path does not exist: {}'.format(self._data_path)

image_path_at

defimage_path_at(self, i):#使用index来从xml_filenames取到filename,生成绝对路径returnself.image_path_from_image_filename(self._image_index[i])

image_path_from_image_filename(类似pascal_voc中的image_path_from_index)

defimage_path_from_image_filename(self, image_filename):image_path = os.path.join(self._data_path,'Images',                                  image_filename + self._image_ext)assertos.path.exists(image_path), \'Path does not exist: {}'.format(image_path)returnimage_path

_load_xml_filenames(类似pascal_voc中的_load_image_set_index)

def_load_xml_filenames(self):#从Annotations文件夹中拿取到bounding box annotation文件名#用来做数据集的索引xml_folder_path = os.path.join(self._data_path,"Annotations")assertos.path.exists(xml_folder_path), \'Path does not exist: {}'.format(xml_folder_path)fordirpath, dirnames, filenamesinos.walk(xml_folder_path):                xml_filenames = [xml_filename.split(".")[0]forxml_filenameinfilenames]returnxml_filenames

gt_roidb

defgt_roidb(self):#Ground-Truth 数据缓存cache_file = os.path.join(self.cache_path, self.name +'_gt_roidb.pkl')ifos.path.exists(cache_file):withopen(cache_file,'rb')asfid:                roidb = cPickle.load(fid)print'{} gt roidb loaded from {}'.format(self.name, cache_file)returnroidb#从xml中获取Ground-Truth数据gt_roidb = [self._load_imagenet_annotation(xml_filename)forxml_filenameinself._image_index]withopen(cache_file,'wb')asfid:            cPickle.dump(gt_roidb, fid, cPickle.HIGHEST_PROTOCOL)print'wrote gt roidb to {}'.format(cache_file)returngt_roidb

rpn_roidb

defrpn_roidb(self):#根据gt_roidb生成rpn_roidb,并进行合并gt_roidb = self.gt_roidb()        rpn_roidb = self._load_rpn_roidb(gt_roidb)        roidb = imdb.merge_roidbs(gt_roidb, rpn_roidb)returnroidb

_load_rpn_roidb

def_load_rpn_roidb(self, gt_roidb):filename = self.config['rpn_file']print'loading {}'.format(filename)assertos.path.exists(filename), \'rpn data not found at: {}'.format(filename)withopen(filename,'rb')asf:            box_list = cPickle.load(f)returnself.create_roidb_from_box_list(box_list, gt_roidb)

_load_imagenet_annotation(类似于pascal_voc中的_load_pascal_annotation)

def_load_imagenet_annotation(self, xml_filename):#从annotation的xml文件中拿取bounding box数据filepath = os.path.join(self._data_path,'Annotations', xml_filename +'.xml')#这里使用了ap,是我写的一个annotation parser,在后面贴出代码#它会返回这个xml文件的wnid, 图像文件名,以及里面包含的注解物体wnid, image_name, objects = ap.parse(filepath)        num_objs = len(objects)        boxes = np.zeros((num_objs,4), dtype=np.uint16)        gt_classes = np.zeros((num_objs), dtype=np.int32)        overlaps = np.zeros((num_objs, self.num_classes), dtype=np.float32)        seg_areas = np.zeros((num_objs), dtype=np.float32)# Load object bounding boxes into a data frame.forix, objinenumerate(objects):            box = obj["box"]            x1 = box['xmin']            y1 = box['ymin']            x2 = box['xmax']            y2 = box['ymax']# 如果这个bounding box并不是我们想要学习的类别,那则跳过# go next if the wnid not exist in declared classestry:                cls = self._class_to_ind[obj["wnid"]]exceptKeyError:print"wnid %s isn't show in given"%obj["wnid"]continueboxes[ix, :] = [x1, y1, x2, y2]            gt_classes[ix] = cls            overlaps[ix, cls] =1.0seg_areas[ix] = (x2 - x1 +1) * (y2 - y1 +1)        overlaps = scipy.sparse.csr_matrix(overlaps)return{'boxes': boxes,'gt_classes': gt_classes,'gt_overlaps': overlaps,'flipped':False,'seg_areas': seg_areas}

annotation_parser.py文件

importosimportxml.dom.minidomdefgetText(node):returnnode.firstChild.nodeValuedefgetWnid(node):returngetText(node.getElementsByTagName("name")[0])defgetImageName(node):returngetText(node.getElementsByTagName("filename")[0])defgetObjects(node):objects = []forobjinnode.getElementsByTagName("object"):        objects.append({"wnid": getText(obj.getElementsByTagName("name")[0]),"box":{"xmin": int(getText(obj.getElementsByTagName("xmin")[0])),"ymin": int(getText(obj.getElementsByTagName("ymin")[0])),"xmax": int(getText(obj.getElementsByTagName("xmax")[0])),"ymax": int(getText(obj.getElementsByTagName("ymax")[0])),            }        })returnobjectsdefparse(filepath):dom = xml.dom.minidom.parse(filepath)    root = dom.documentElement    image_name = getImageName(root)    wnid = getWnid(root)    objects = getObjects(root)returnwnid, image_name, objects

则对数据结构的要求是:

|---data|---imagenet|---Annotations|---n03147509|---n03147509_*.xml|---...|---n04272054|---n04272054_*.xml|---...|---Images|---n03147508_*.JPEG|---...|---n04272054_*.JPEG|---...

同时我在github上也提供了draw方法,可以用来将bounding box画于Image文件上,用来甄别该annotation的正确性

训练

这样,我们的ImageNet类则是生成好了,下面我们则可以训练我们的数据,但是在开始之前,还有一件事情,那就是修改prototxt中的与类别数目有关的值,我将models/pascal_voc拷贝到了models/imagenet进行修改,比如我想要训练ZF,如果使用的是train_faster_rcnn_alt_opt.py,则需要修改models/imagenet/ZF/faster_rcnn_alt_opt/下的所有pt文件里的内容,用如下的法则去替换:

//num为类别的个数input-data->num_classes = numclass_score->num_output = numbbox_pred->num_output  = num*4

我这里使用train_faster_rcnn_alt_opt.py进行的训练,这样的话则需要把添加的models/imagenet作为可选项

//pt_type 则是添加的选择项,默认使用psacal_voc的models./tools/train_faster_rcnn_alt_opt.py--gpu 0 \--net_name ZF \--weights data/imagenet_models/ZF.v2.caffemodel[optional] \--imdb imagenet \--cfg experiments/cfgs/faster_rcnn_alt_opt.yml \--pt_type imagenet

识别

这里我们则需要使用刚训练出来的模型进行识别

#就像demo.py一样,但是使用训练的models,我创建了tools/classify.py来单独识别prototxt = os.path.join(cfg.ROOT_DIR, 'models/imagenet', NETS[args.demo_net][0], 'faster_rcnn_alt_opt', 'faster_rcnn_test.pt')caffemodel = os.path.join(cfg.ROOT_DIR, 'output/faster_rcnn_alt_opt/imagenet/'+ NETS[args.demo_net][0] +'_faster_rcnn_final.caffemodel')

同样,在识别前我们要对识别方法里的Classes进行修改,修改成你自己训练的类别后

执行

./tools/classify.py--net zf

则可对data/demo下的图片文件使用训练的zf网络进行识别

Have fun

你可能感兴趣的:([转载]使用ImageNet在faster-rcnn上训练自己的分类网络)