affectnet.py 3.74 KB
Newer Older
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
import argparse
import logging
import os
from shutil import copyfile

import pandas as pd

parser = argparse.ArgumentParser(description='AffectNet PrePipeline')
parser.add_argument('csv', action='store', help='the .csv file to process')
parser.add_argument('source', action='store', help='image source folder')
parser.add_argument('-0', action='append_const', dest='emotions', const='neutral', help='neutral')
parser.add_argument('-1', action='append_const', dest='emotions', const='happy', help='happy')
parser.add_argument('-2', action='append_const', dest='emotions', const='sadness', help='sadness')
parser.add_argument('-3', action='append_const', dest='emotions', const='surprise', help='surprise')
parser.add_argument('-4', action='append_const', dest='emotions', const='fear', help='fear')
parser.add_argument('-5', action='append_const', dest='emotions', const='disgust', help='disgust')
parser.add_argument('-6', action='append_const', dest='emotions', const='anger', help='anger')
parser.add_argument('-7', action='append_const', dest='emotions', const='contempt', help='contempt')
parser.add_argument('-8', action='append_const', dest='emotions', const='none', help='none')
parser.add_argument('-9', action='append_const', dest='emotions', const='uncertain', help='uncertain')
parser.add_argument('-d', action='store', dest='destination', default='sorted_set', help='destination to sorted set')
parser.add_argument('-l', '--limit', action='store', type=int, dest='limit', default=-1, help='image limit per emotion')
parser.add_argument('-v', '--valence', action='store', default='-2', dest='valence', type=float, help='min valence')
parser.add_argument('-a', '--arousal', action='store', default='-2', dest='arousal', type=float, help='min arousal')
arguments = parser.parse_args()
logging.debug(arguments)

tihmels's avatar
tihmels committed
28
29
emotidict = {0: 'neutral', 1: 'happy', 2: 'sadness', 3: 'surprise', 4: 'fear', 5: 'disgust', 6: 'anger', 7: 'contempt',
             8: 'none', 9: 'uncertain', 10: 'noface'}
30

tihmels's avatar
tihmels committed
31
emoticounter = {}
32
for e in arguments.emotions:
tihmels's avatar
tihmels committed
33
    emoticounter[e] = 0
34
35
36
37
38
39

curdir = os.path.abspath(os.path.dirname(__file__))
csv = os.path.join(curdir, arguments.csv)
source = os.path.join(curdir, arguments.source)
destination = os.path.join(curdir, arguments.destination)

tihmels's avatar
tihmels committed
40

41
42
43
44
45
46
47
48
49
50
51
52
def generate_records(csvfile):
    data = pd.read_csv(csvfile, delimiter=',', dtype='a')

    imagepaths = pd.np.array(data['subDirectory_filePath'], pd.np.str)
    labels = pd.np.array(data['expression'], pd.np.int)
    valence = pd.np.array(data['valence'], pd.np.float)
    arousal = pd.np.array(data['arousal'], pd.np.float)

    for e in arguments.emotions:
        dest = os.path.join(destination, e)
        if not os.path.exists(dest):
            os.makedirs(dest)
tihmels's avatar
tihmels committed
53
        else:
tihmels's avatar
tihmels committed
54
            emoticounter[e] = len(os.listdir(dest))
55
56
57
58

    data = zip(labels, imagepaths, valence, arousal)

    for d in data:
tihmels's avatar
tihmels committed
59
        emotion = emotidict.get(d[0])
tihmels's avatar
tihmels committed
60

tihmels's avatar
tihmels committed
61
62
        if emotion not in arguments.emotions or d[2] < arguments.valence or d[3] < arguments.arousal or emoticounter[
            emotion] >= arguments.limit:
tihmels's avatar
tihmels committed
63
64
            continue

65
66
67
68
69
70
71
        img = d[1]
        pathfrom = os.path.join(source, img)

        if not os.path.exists(pathfrom):
            logging.info(img + ' not found')
            continue

tihmels's avatar
tihmels committed
72
        folder = os.path.join(destination, emotion)
tihmels's avatar
tihmels committed
73
        filePath = os.path.join(folder, 'V' + str(round(d[2], 1)) + 'A' + str(round(d[3], 1)) + '_' + img.split('/')[1])
74

tihmels's avatar
tihmels committed
75
76
77
78
        if not os.path.isfile(filePath):
            copyfile(pathfrom, filePath)
            logging.info('wrote file ' + filePath)
            emoticounter[emotion] += 1
tihmels's avatar
tihmels committed
79

tihmels's avatar
tihmels committed
80
81
        if all(count >= arguments.limit for count in emoticounter.values()):
            break
82
83
84
85
86
87


if __name__ == '__main__':
    generate_records(csv)

    logging.info('finish')
tihmels's avatar
tihmels committed
88
    logging.info(emoticounter)