def setup_net(self): caffe.set_mode_gpu() caffe.set_device(0) netFiles = cfg.get_caffe_net_files(self.prms_.netName) self.net_ = caffe.Net(netFiles.deployFile, netFiles.netFile, caffe.TEST) self.cls_ = cfg.dataset2classnames(self.prms_.trainDataSet)
def __init__(self, prms={}): # define location for output file and image folder try: self.outFile = prms['outFile'] except KeyError: outFolder = '/mnt/Ext/training_data/rcnn' dateExt = cuu.get_datetime_ext() self.outFile = osp.join(outFolder, dateExt, 'rcnn_' + dateExt + '.txt') os.system('mkdir -p ' + osp.dirname(self.outFile)) prms['outFile'] = self.outFile try: self.imageFolder = prms['imageFolder'] except KeyError: self.imageFolder = osp.join(osp.dirname(self.outFile), 'imageFolder') os.system('mkdir -p ' + self.imageFolder) prms['imageFolder'] = self.imageFolder # define dataset used for pretraining self.pretrainDataSet = cfg.get_rcnn_prms(**prms)['trainDataSet'] # match available labels with available classes from the pretraining dataset clsNames = cfg.dataset2classnames(self.pretrainDataSet) labeledObjs = cfg.get_label_types(category='object') self.clsLookup = {} for obj in labeledObjs: try: self.clsLookup[obj] = clsNames.index(cfg.label2dataset(obj, dataset=self.pretrainDataSet)) except ValueError as e: print e ch.ChainObject.__init__(self, prms)