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)
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)