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')