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