affectnet.py 3.87 KB
Newer Older
1
2
3
4
5
6
7
8
9
import argparse
import logging
import os
import shutil
import sys
from shutil import copyfile

import pandas as pd

tihmels's avatar
tihmels committed
10
logfile = 'affectnet.log'
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59

# Erstellt und konfiguriert die Logdatei
logging.basicConfig(level=logging.NOTSET, format='%(asctime)s %(levelname)-8s %(message)s',
                    datefmt='%m-%d %H:%M',
                    filename=logfile)

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)

emodict = {0: 'neutral', 1: 'happy', 2: 'sadness', 3: 'surprise', 4: 'fear', 5: 'disgust', 6: 'anger', 7: 'contempt', 8: 'none', 9: 'uncertain', 10: 'noface'}

counter = {}
for e in arguments.emotions:
    counter[e] = 0

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)

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
60
61
        else:
            counter[e] = len(os.listdir(dest))
62
63
64
65

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

    for d in data:
tihmels's avatar
tihmels committed
66
67
68
69
70
71
72
73
        if all(count >= arguments.limit for count in counter.values()):
            break

        emotion = emodict.get(d[0])

        if emotion not in arguments.emotions or d[2] < arguments.valence or d[3] < arguments.arousal or counter[emotion] >= arguments.limit:
            continue

74
75
76
77
78
79
80
81
82
83
        img = d[1]
        pathfrom = os.path.join(source, img)

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

        parent = os.path.join(destination, emotion)

        pathto = os.path.join(parent, img.split('/')[1])
tihmels's avatar
tihmels committed
84
85
86
87
88

        if not os.path.isfile(pathto):
            copyfile(pathfrom, pathto)
            logging.info('wrote file ' + pathto)
            counter[emotion] += 1
89
90
91
92
93
94
95


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

    logging.info('finish')
    logging.info(counter)