import argparse import os from PIL import Image import time import json parser = argparse.ArgumentParser(description=""" Frames are split into segments so SMALL_FACES actually detects _really_ small faces. This requires enabling --segments on opencv-webcam-demo (so it detects/writes segment%06d files) split_and_merge_output.py also remerges these segment%06.json files into a frame%06d.json file. """) parser.add_argument('--frameOutput', '-o', required=True, help='directory to look for frames & json') parser.add_argument('--segmentDir', required=True, help='directory to write segfor frames & json') args = parser.parse_args() if not os.path.exists(args.frameOutput): raise Exception(args.frameOutput + " does not exist") if not os.path.exists(args.segmentDir): raise Exception(args.segmentDir + " does not exist") def getSectionsForSize(size): """ Get sections based on image size, in a tuple format (width, height) """ return [ ( 0, 0, int(size[0]/1.85), int(size[1]/1.85) ), ( int(size[0]/2.1), 0, size[0], int(size[1]/1.85) ), ( 0, int(size[1]/2.1), int(size[0]/1.85), size[1] ), ( int(size[0]/2.1), int(size[1]/2.1), size[0], size[1] ), ] nextInNr = 1 segmentCount = 4 # cannot really be changed for now nextOutNr = segmentCount while True: nextJpg = os.path.join(args.frameOutput, "frame%06d.jpg" % nextInNr) nextInJson = os.path.join(args.frameOutput, "frame%06d.json" % nextInNr) while os.path.exists(nextJpg) and not os.path.exists(nextInJson): print("(SPLIT) Found: {}".format(nextJpg)) img = Image.open(nextJpg) segments = getSectionsForSize(img.size) for i, segment in enumerate(segments): segmentImg = img.crop(segment) segmentNr = (nextInNr-1) * segmentCount + i + 1 # start at 1, increment from there segmentFilename = os.path.join(args.segmentDir, "segment%06d.jpg" % segmentNr) segmentImg.save(segmentFilename, quality=90) print("(SPLIT) Created: {}".format(segmentFilename)) nextInNr += 1 nextJpg = os.path.join(args.frameOutput, "frame%06d.jpg" % nextInNr) nextInJson = os.path.join(args.frameOutput, "frame%06d.jpg" % nextInNr) nextJson = os.path.join(args.segmentDir, "segment%06d.json" % nextOutNr) while os.path.exists(nextJson): print("(MERGE) Found: {}".format(nextJson)) frameOutNr = (nextOutNr - 1) / segmentCount + 1 # 1,2,3,4 -> 1; 5,6,7,8 -> 2 (uses trick to divide ints: 7/4=1) img = Image.open( os.path.join(args.frameOutput, "frame%06d.jpg" % frameOutNr) ) segments = getSectionsForSize(img.size) creationTime = img._getexif()[36867] faces = [] times = [] segmentImages = [] for i, segment in enumerate(segments): segmentNr = nextOutNr - segmentCount + i + 1 # nextOutNr: 4 => 1,2,3,4 segmentImages.append(os.path.join(args.segmentDir, "segment%06d.jpg" % segmentNr)) with open( os.path.join(args.segmentDir, "segment%06d.json" % segmentNr) ) as fp: j = json.load(fp) for f in j['faces']: f['rect']['x'] += segment[0] f['rect']['y'] += segment[1] f['id'] = "%d%02d" % ((i+1), f['id']) faces.append(f) times.append(j['t']) data = { 'nr': int(frameOutNr), 't': creationTime, 'faces': faces, 'sectionTimes': times, } jsonOutFilename = os.path.join(args.frameOutput, "frame%06d.json" % frameOutNr) with open(jsonOutFilename, 'w') as fp: json.dump(data, fp) print("(MERGE) Wrote: {}".format(jsonOutFilename)) print("(MERGE) Cleanup segments for frame {}".format(frameOutNr)) for imgName in segmentImages: print("(MERGE) \t Unlink {}".format(imgName)) os.unlink(imgName) nextOutNr += segmentCount nextJson = os.path.join(args.segmentDir, "segment%06d.json" % nextOutNr) time.sleep(0.2)