105 lines
3.8 KiB
Python
105 lines
3.8 KiB
Python
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):
|
|
if 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))
|
|
else:
|
|
print("(SPLIT) Skip: {}".format(nextJpg))
|
|
|
|
nextInNr += 1
|
|
nextJpg = os.path.join(args.frameOutput, "frame%06d.jpg" % nextInNr)
|
|
nextInJson = os.path.join(args.frameOutput, "frame%06d.json" % nextInNr)
|
|
|
|
continue
|
|
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:
|
|
if os.path.exists(imgName):
|
|
print("(MERGE) \t Unlink {}".format(imgName))
|
|
os.unlink(imgName)
|
|
else:
|
|
print("(MERGE) \t Already unlinked {}".format(imgName))
|
|
|
|
nextOutNr += segmentCount
|
|
nextJson = os.path.join(args.segmentDir, "segment%06d.json" % nextOutNr)
|
|
|
|
time.sleep(0.2)
|