process_model.py 2.66 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
args = sys.argv
tihmels's avatar
tihmels committed
12
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
def image_preprocessing(item):
    image = cv2.imread(item)  # open image
    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)  # convert to grayscale
tihmels's avatar
tihmels committed
24
    return gray
Arne Gerdes's avatar
Arne Gerdes committed
25
26
27
28
29
30
31

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

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

tihmels's avatar
tihmels committed
44

Arne Gerdes's avatar
Arne Gerdes committed
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
    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
65
66
67
if len(args) > 1:
    tags = ', '.join(args[1:])
    logging.debug(tags.upper())
tihmels's avatar
tihmels committed
68
69
70

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

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

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

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