face-sentiment.py

import csv
import io
import os

from google.cloud import vision


def detect_faces(path):
    client = vision.ImageAnnotatorClient()

    with io.open(path, 'rb') as image_file:
        content = image_file.read()
    image = vision.types.Image(content=content)

    response = client.face_detection(image=image)
    faces = response.face_annotations

    likelihood_name = ('UNKNOWN', 'VERY_UNLIKELY', 'UNLIKELY', 'POSSIBLE',
                       'LIKELY', 'VERY_LIKELY')
    joy, sorrow, anger, surprise = 0, 0, 0, 0

    for face in faces:
        if face.joy_likelihood > 3:
            joy += 1
        if face.sorrow_likelihood > 3:
            sorrow += 1
        if face.anger_likelihood > 3:
            anger += 1
        if face.surprise_likelihood > 3:
            surprise += 1

    return [path.split('/')[-1], joy, sorrow, anger, surprise]


if __name__ == '__main__':
    filenames = os.listdir('frames')
    filenames.sort(key=lambda x: int(''.join(filter(str.isdigit, x))))
    csv_header = ['n_frame', 'joy', 'sorrow', 'anger', 'surprise']
    csv_results = [csv_header]
    for frame in filenames[823:]:
        result = detect_faces('frames/{}'.format(frame))
        print(result)
        csv_results.append(result)

    csv_file = open('results.csv', 'w')
    with csv_file:
        writer = csv.writer(csv_file)
        writer.writerows(csv_results)