specimens-of-composite-port.../portrait_compositor.py

273 lines
7.2 KiB
Python

import picamera
import io, os
import datetime
import csv
from subprocess import Popen, PIPE
from PIL import Image
import numpy as np
import cPickle as pickle
import requests
import time
import thread
from websocket import create_connection
import logging
import json
camera = picamera.PiCamera()
camera.rotation = 180
camera.resolution = (1920, 1080)
# camera.resolution = (1280, 720)
outputResolution = (1000, 1000)
# the binary genders as outputted by Affectiva
genders = ['male', 'female', 'unknown']
perspectives = ['side', 'front']
gender_perspectives = [g+"_"+p for p in perspectives for g in genders]
curdir = os.path.dirname(os.path.abspath(__file__))
tmpimage = '/tmp/piimage.jpg'
tmpimageResults = '/tmp/piimage.csv'
cmd = [
os.path.join(curdir, 'build/video-demo/video-demo'),
'--input', tmpimage,
'--data', os.path.join(curdir, 'affdex-sdk/data'),
'--draw', '0',
'--numFaces', '20',
]
# without these vars video-demo yields a segmentation fault
environment = {
'LC_LANG': 'en_GB.UTF-8',
'LD_PRELOAD': '/usr/lib/arm-linux-gnueabihf/libopencv_core.so.2.4',
}
def updateStats(type, name, count, image_filename):
params = {
'type': type,
'name': name,
'time': int(time.time()),
'case_count': int(count),
}
try:
ws = create_connection("ws://localhost:8888/ws")
js = json.dumps({
'type': type,
'name': name,
'img_src': os.path.basename(image_filename),
'case_count': int(count),
})
ws.send(js)
except Exception as e:
raise
url = 'https://artstats.rubenvandeven.com/composites/views.php'
if count % 10 == 0:
# only send every one in x image, so that the server never can
# retrace _exact_ faces by comparing the sent frames.
with open(image_filename) as fp:
print('send request including image')
r = requests.post(
url ,
files={'image': fp},
params=params
)
else:
print('send request')
r = requests.post(
url,
params=params
)
class CompositeImage:
def __init__(self, name, resolution):
self.name = name
self.count = 0
self.resolution = resolution
self.image = np.zeros((resolution[0],resolution[1],3))
# use state to determine whether a save is necessary
self.state_dirty = True
def addFace(self, img):
img_a = np.array(img.resize(self.resolution))
self.count += 1
self.image = (self.image * (self.count - 1)/float(self.count) + img_a / float(self.count))
self.state_dirty = True
def restore(self, i, dir):
'''
Restore from pickle nr
'''
self.count = i
name = self.get_frame_filename(self.count)
img_file = os.path.join(dir, name)
print("\trestore {}".format(img_file))
self.image = np.array(Image.open(img_file)).astype('float64')
self.state_dirty = False
def get_image(self):
return Image.fromarray(self.image.astype('uint8'),'RGB')
def get_frame_filename(self, i):
return "{}-{}x{}-{}.png".format(self.name, self.resolution[0], self.resolution[1], i)
def save_image(self, dir):
if self.state_dirty is False:
return
name = self.get_frame_filename(self.count)
filename = os.path.join(dir, name)
self.get_image().save(filename)
thread.start_new_thread( updateStats, ('gender', self.name, self.count, filename) )
self.state_dirty = False
class CompositeCollection:
"""
Store/save the composite images
"""
def __init__(self, names, size, target_dir = None):
self.id = "{}-{}x{}".format("-".join(names), size[0], size[1])
self.names = names
self.size = size
self.target_dir = os.path.dirname(os.path.abspath(__file__)) if target_dir is None else target_dir
self.load()
def get_pickle_filename(self):
return os.path.join(self.target_dir, self.id + ".p")
def load(self):
pickle_file_name = self.get_pickle_filename()
# if os.path.exists(pickle_file_name):
composites = {}
try:
with open( pickle_file_name, "rb" ) as fp:
data = pickle.load( fp )
for name in data['c']:
composites[name] = CompositeImage(name, self.size)
composites[name].restore( data['c'][name], self.target_dir)
except Exception as e:
print("Create new composite", e)
for name in self.names:
composites[name] = CompositeImage(name, self.size)
self.composites = composites
def save(self):
data = { 'size' : self.size, 'c': {} }
for name in self.composites:
data['c'][name] = self.composites[name].count
with open( self.get_pickle_filename(), "wb" ) as fp:
print("Save", data)
pickle.dump( data, fp )
def save_img(self, name):
self.get(name).save_image(self.target_dir)
def get_as_percentages(self, precision = 3):
total = sum([c.count for c in self.composites])
percentages = {}
if total < 1:
# assert: in the beginning, we were all made equal
for c in self.composites:
percentages[c.name] = round(100 / len(self.composites), precision)
else:
for c in self.composites:
percentages[c.name] = round(100 * (c.count / total), precision)
return percentages
def get(self, name):
return self.composites[name]
def clean(self):
for name in self.names:
c = self.get(name)
start = max(0, c.count - 10)
end = max(0, c.count - 5)
for i in range(start, end):
filename = os.path.join(self.target_dir, c.get_frame_filename(i))
if os.path.exists(filename):
print("Clean {}".format(filename))
os.unlink(filename)
def append_face(row, image, composites):
# degrees to distinguish side (as we will never be able to use 90 :-( )
suffix = 'side' if abs(float(row['yaw'])) > 20 else 'front'
# print('yaw:', float(row['yaw']))
name = "{}_{}".format(row['gender'], suffix)
if name not in composites.names:
return
composite = composites.get(name)
# TODO: matrix transform the image, to skew the face into being a flat-ish surface
# This might yield less blurry composites
# crop image, bt keep it bigger than the found face
grow_x = .2 # in every direction, so .2 becomes 1.4 * width
grow_y = grow_x
face_w = int(row['width'])
face_h = int(row['height'])
face_x = int(row['x'])
face_y = int(row['y'])
# we go square:
size_x = max(face_w, face_h) * (1 + grow_x * 2)
size_y = size_x
dx = (face_w - size_x) / 2
dy = (face_h - size_y) / 2
# PIL.Image handles cropping outside the canvas by filling with black/transparent
x = face_x + dx
y = face_y + dy
print('crop')
i = image.crop((x,y, x + size_x, y + size_y))
if suffix == 'side' and float(row['yaw']) < 0:
print('\tflip')
i = i.transpose(Image.FLIP_LEFT_RIGHT)
print('add')
composite.addFace(i)
print('added')
composites = CompositeCollection(gender_perspectives, outputResolution, os.path.join(curdir, 'output'))
while True:
start = datetime.datetime.utcnow()
# stream = io.BytesIO()
camera.capture(tmpimage, format='jpeg')
process = Popen(cmd, env=environment)
process.wait()
img = Image.open(tmpimage)
os.unlink(tmpimage)
with open(tmpimageResults) as csvfile:
print("open csv")
data = csv.DictReader(csvfile)
faces = 0
for row in data:
if row['faceId'] == 'nan':
# not a valid face
continue
faces += 1
print("append face")
append_face(row, img, composites)
if faces > 0:
print("save :-)")
for name in composites.names:
print("\tsave img '{}'".format(name))
c = composites.save_img(name)
# save pickle after images, so they can be restored
composites.save()
composites.clean()
# TODO: trigger output update