affectnet.py 4.75 KB
Newer Older
tihmels's avatar
tihmels committed
1
2
3
4
5
6
7
8
9
# ***************************************************
# * Copyright (C) 2010-2011 Tjado Ihmels <tjado.ihmels@uni-oldenburg.de>
# *
# * This file is part of whiteboard-project-matcher.
# *
# * whiteboard-project-matcher can not be copied and/or distributed without the express
# * permission of Tjado Ihmels
# ***************************************************/

10
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
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
37
38
39
# the numbers define the particular emotions in the .csv file
emotiondict = {0: 'neutral', 1: 'happy', 2: 'sadness', 3: 'surprise', 4: 'fear', 5: 'disgust', 6: 'anger', 7: 'contempt',
               8: 'none', 9: 'uncertain', 10: 'noface'}
40

tihmels's avatar
tihmels committed
41
# count found emotions
tihmels's avatar
tihmels committed
42
43
44
emotioncounter = {}
for emotion in arguments.emotions:
    emotioncounter[emotion] = 0
45

tihmels's avatar
tihmels committed
46
47
48
49
50
# define some useful paths
currentdir = os.path.abspath(os.path.dirname(__file__))
csv = os.path.join(currentdir, arguments.csv)
source = os.path.join(currentdir, arguments.source)
destination = os.path.join(currentdir, arguments.destination)
tihmels's avatar
tihmels committed
51

52
def generate_records(csvfile):
tihmels's avatar
tihmels committed
53
    # read the .csv file
54
55
    data = pd.read_csv(csvfile, delimiter=',', dtype='a')

tihmels's avatar
tihmels committed
56
    # get the necessary columns
57
58
59
60
61
    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)

tihmels's avatar
tihmels committed
62
63
64
65
66
    for emotion in arguments.emotions:
        # join the destination path with the emotion identifier and create it if it isn't..
        dest_subfolder = os.path.join(destination, emotion)
        if not os.path.exists(dest_subfolder):
            os.makedirs(dest_subfolder)
tihmels's avatar
tihmels committed
67
        else:
tihmels's avatar
tihmels committed
68
69
            # .. otherwise, count images
            emotioncounter[emotion] = len(os.listdir(dest_subfolder))
70

tihmels's avatar
tihmels committed
71
72
    # wrap tupels for iterating
    csvrows = zip(labels, imagepaths, valence, arousal)
73

tihmels's avatar
tihmels committed
74
75
76
    for c in csvrows:
        # get emotion
        emotion = emotiondict.get(c[0])
tihmels's avatar
tihmels committed
77

tihmels's avatar
tihmels committed
78
79
        # check if the row fits the argument conditions
        if emotion not in arguments.emotions or c[2] < arguments.valence or c[3] < arguments.arousal or emotioncounter[
tihmels's avatar
tihmels committed
80
            emotion] >= arguments.limit:
tihmels's avatar
tihmels committed
81
82
            continue

tihmels's avatar
tihmels committed
83
84
        # the image path
        img = c[1]
85
86
        pathfrom = os.path.join(source, img)

tihmels's avatar
tihmels committed
87
        # check if the respective file exists
88
89
90
91
        if not os.path.exists(pathfrom):
            logging.info(img + ' not found')
            continue

tihmels's avatar
tihmels committed
92
93
94
        dest_subfolder = os.path.join(destination, emotion)
        # join destination path with filename
        filePath = os.path.join(dest_subfolder, img.split('/')[1])
95

tihmels's avatar
tihmels committed
96
        # if not exists, copy file to destination
tihmels's avatar
tihmels committed
97
98
99
        if not os.path.isfile(filePath):
            copyfile(pathfrom, filePath)
            logging.info('wrote file ' + filePath)
tihmels's avatar
tihmels committed
100
            emotioncounter[emotion] += 1
tihmels's avatar
tihmels committed
101

tihmels's avatar
tihmels committed
102
103
        # break if enough files were copied
        if all(count >= arguments.limit for count in emotioncounter.values()):
tihmels's avatar
tihmels committed
104
            break
105
106
107
108
109


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

tihmels's avatar
tihmels committed
110
111
    print('finish')
    print(emotioncounter)