コード例 #1
0
ファイル: jpkg_chains.py プロジェクト: pulkitag/chainer
    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)
コード例 #2
0
ファイル: vis_chains.py プロジェクト: pulkitag/chainer
 def produce(self, ip):
         # select the n best and n worst
         ip.sort(key=lambda x: x[1])
         if len(ip) >= 2*self._n:
                 sfs = ip[:self._n] + ip[-self._n:]
         else:
                 sfs = ip
         # form videos from the sets of images
         dtExt = cu.get_datetime_ext()
         for imageDir, predictions, mAP in sfs: 
                 # form input, output paths
                 if np.isnan(mAP):
                         mAP = 0.0
                 vidName  = dtExt + '_' + 'mAP-' + str(int(mAP*1000)) + '.mp4'
                 vidPath  = os.path.join(self.outDir, vidName) 
                 tmpDir   = os.path.join(self.outDir, 'tmp')
                 while os.path.exists(tmpDir):
                         tmpDir += '1'
                 os.mkdir(tmpDir)
                 # write bounding boxes to image
                 for dirpath, dirnames, filenames in os.walk(imageDir):
                         for filename in filenames:
                                 im          = scm.imread(os.path.join(dirpath, filename))
                                 imBoxer     = MakeImBBox()
                                 frameNum    = int(filename.split('.')[0]) + 1
                                 try:
                                         framePreds  = [p[2] for p in predictions[frameNum]]
                                 except IndexError:
                                         framePreds  = []
                                 if framePreds:
                                         im  = imBoxer.produce((im, framePreds))
                                 scm.imsave(os.path.join(tmpDir, '{0:04d}.jpg'.format(frameNum)), im)
                 # make video
                 ffCmd = 'ffmpeg -y -framerate 5 -i ' + \
                         os.path.join(tmpDir, '%04d.jpg') + ' -c:v mpeg4 ' + vidPath
                 rmCmd = 'rm -r ' + tmpDir
                 sp.Popen(ffCmd + '; ' + rmCmd + '; exit 0;', shell=True)