affectnet.py 4.05 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 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 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99
import argparse
import logging
import os
import shutil
import sys
from shutil import copyfile

import pandas as pd

logfile = '../logs/affectnet.log'

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

if os.path.exists(destination):
    delete = input(destination + ' existiert. Durch diesen Prozess werden die Daten gelöscht. Fortfahren (y/n): ')
    if delete == 'y':
        shutil.rmtree(destination)
    else:
        logging.info('process aborted')
        sys.exit()

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)

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

    for d in data:
        img = d[1]
        pathfrom = os.path.join(source, img)

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

        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] < float(arguments.valence) or d[3] < float(arguments.arousal) or counter[emotion] == arguments.limit:
            continue

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

        pathto = os.path.join(parent, img.split('/')[1])
        copyfile(pathfrom, pathto)
        logging.info('wrote file ' + pathto)
        counter[emotion] += 1


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

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