process_model.py 2.7 KB
Newer Older
Arne Gerdes's avatar
Arne Gerdes committed
1 2 3 4
import cv2
import glob
import random
import numpy as np
tihmels's avatar
tihmels committed
5
import sys
tihmels's avatar
tihmels committed
6
import logging
Arne Gerdes's avatar
Arne Gerdes committed
7

tihmels's avatar
tihmels committed
8
logging.basicConfig(level=logging.DEBUG, format='%(asctime)s %(levelname)-8s %(message)s',
tihmels's avatar
tihmels committed
9 10
                    datefmt='%m-%d %H:%M',
                    filename='logs/process_model.log')
tihmels's avatar
tihmels committed
11 12
tags = sys.argv
logging.debug('Fisherface training initialized')
Arne Gerdes's avatar
Arne Gerdes committed
13

tihmels's avatar
tihmels committed
14
def get_files_from_emotion(emotion):  # Define function to get file list, randomly shuffle it and split 80/20
tihmels's avatar
tihmels committed
15
    files = glob.glob('basis_data/dataset/%s/*' % emotion)
Arne Gerdes's avatar
Arne Gerdes committed
16 17 18 19 20
    random.shuffle(files)
    training = files[:int(len(files) * 0.8)]  # get first 80% of file list
    prediction = files[-int(len(files) * 0.2):]  # get last 20% of file list
    return training, prediction

tihmels's avatar
tihmels committed
21 22 23 24 25
def image_preprocessing(item):
    image = cv2.imread(item)  # open image
    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)  # convert to grayscale
    equalized = cv2.equalizeHist(gray)
    return equalized
Arne Gerdes's avatar
Arne Gerdes committed
26 27 28 29 30 31 32

def make_sets():
    training_data = []
    training_labels = []
    prediction_data = []
    prediction_labels = []
    for emotion in emotions:
tihmels's avatar
tihmels committed
33
        training, prediction = get_files_from_emotion(emotion)
Arne Gerdes's avatar
Arne Gerdes committed
34 35
        # Append data to training and prediction list, and generate labels 0-7
        for item in training:
tihmels's avatar
tihmels committed
36 37
            img = image_preprocessing(item)
            training_data.append(img)  # append image array to training data list
Arne Gerdes's avatar
Arne Gerdes committed
38 39 40
            training_labels.append(emotions.index(emotion))

        for item in prediction:  # repeat above process for prediction set
tihmels's avatar
tihmels committed
41 42
            img = image_preprocessing(item)
            prediction_data.append(img)
Arne Gerdes's avatar
Arne Gerdes committed
43 44
            prediction_labels.append(emotions.index(emotion))

tihmels's avatar
tihmels committed
45

Arne Gerdes's avatar
Arne Gerdes committed
46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65
    return training_data, training_labels, prediction_data, prediction_labels

def run_recognizer():
    training_data, training_labels, prediction_data, prediction_labels = make_sets()

    fishface.train(training_data, np.asarray(training_labels))

    cnt = 0
    correct = 0
    incorrect = 0
    for image in prediction_data:
        pred, conf = fishface.predict(image)
        if pred == prediction_labels[cnt]:
            correct += 1
            cnt += 1
        else:
            incorrect += 1
            cnt += 1
    return ((100 * correct) / (correct + incorrect))

tihmels's avatar
tihmels committed
66 67 68 69 70 71
if len(tags) > 1:
    message = ', '.join(tags[1:])
    logging.debug(message.upper())

emotions = ["anger", "disgust", "happy", "neutral", "surprise"]  # Emotion list
fishface = cv2.face.FisherFaceRecognizer_create()
Arne Gerdes's avatar
Arne Gerdes committed
72 73 74

# Now run it
metascore = []
tihmels's avatar
tihmels committed
75

Arne Gerdes's avatar
Arne Gerdes committed
76 77
for i in range(0, 10):
    correct = run_recognizer()
tihmels's avatar
tihmels committed
78
    logging.debug("{} : {}%".format(i, int(correct)))
Arne Gerdes's avatar
Arne Gerdes committed
79 80
    metascore.append(correct)

tihmels's avatar
tihmels committed
81
logging.debug("{} iterations - {}% average\n".format(len(metascore), np.mean(metascore)))
tihmels's avatar
tihmels committed
82
fishface.write('basis_data/models/detection_model.xml')