Commit 666f8d03 authored by tihmels's avatar tihmels

AffectNet Pipeline

parent 5af0c9a4
import argparse import argparse
import logging import logging
import os import os
import shutil
import sys
from shutil import copyfile from shutil import copyfile
import pandas as pd import pandas as pd
...@@ -34,17 +32,19 @@ parser.add_argument('-a', '--arousal', action='store', default='-2', dest='arous ...@@ -34,17 +32,19 @@ parser.add_argument('-a', '--arousal', action='store', default='-2', dest='arous
arguments = parser.parse_args() arguments = parser.parse_args()
logging.debug(arguments) 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'} emotidict = {0: 'neutral', 1: 'happy', 2: 'sadness', 3: 'surprise', 4: 'fear', 5: 'disgust', 6: 'anger', 7: 'contempt',
8: 'none', 9: 'uncertain', 10: 'noface'}
counter = {} emoticounter = {}
for e in arguments.emotions: for e in arguments.emotions:
counter[e] = 0 emoticounter[e] = 0
curdir = os.path.abspath(os.path.dirname(__file__)) curdir = os.path.abspath(os.path.dirname(__file__))
csv = os.path.join(curdir, arguments.csv) csv = os.path.join(curdir, arguments.csv)
source = os.path.join(curdir, arguments.source) source = os.path.join(curdir, arguments.source)
destination = os.path.join(curdir, arguments.destination) destination = os.path.join(curdir, arguments.destination)
def generate_records(csvfile): def generate_records(csvfile):
data = pd.read_csv(csvfile, delimiter=',', dtype='a') data = pd.read_csv(csvfile, delimiter=',', dtype='a')
...@@ -58,17 +58,15 @@ def generate_records(csvfile): ...@@ -58,17 +58,15 @@ def generate_records(csvfile):
if not os.path.exists(dest): if not os.path.exists(dest):
os.makedirs(dest) os.makedirs(dest)
else: else:
counter[e] = len(os.listdir(dest)) emoticounter[e] = len(os.listdir(dest))
data = zip(labels, imagepaths, valence, arousal) data = zip(labels, imagepaths, valence, arousal)
for d in data: for d in data:
if all(count >= arguments.limit for count in counter.values()): emotion = emotidict.get(d[0])
break
emotion = emodict.get(d[0]) if emotion not in arguments.emotions or d[2] < arguments.valence or d[3] < arguments.arousal or emoticounter[
emotion] >= arguments.limit:
if emotion not in arguments.emotions or d[2] < arguments.valence or d[3] < arguments.arousal or counter[emotion] >= arguments.limit:
continue continue
img = d[1] img = d[1]
...@@ -78,18 +76,20 @@ def generate_records(csvfile): ...@@ -78,18 +76,20 @@ def generate_records(csvfile):
logging.info(img + ' not found') logging.info(img + ' not found')
continue continue
parent = os.path.join(destination, emotion) folder = os.path.join(destination, emotion)
filePath = os.path.join(folder, img.split('/')[1])
pathto = os.path.join(parent, img.split('/')[1]) if not os.path.isfile(filePath):
copyfile(pathfrom, filePath)
logging.info('wrote file ' + filePath)
emoticounter[emotion] += 1
if not os.path.isfile(pathto): if all(count >= arguments.limit for count in emoticounter.values()):
copyfile(pathfrom, pathto) break
logging.info('wrote file ' + pathto)
counter[emotion] += 1
if __name__ == '__main__': if __name__ == '__main__':
generate_records(csv) generate_records(csv)
logging.info('finish') logging.info('finish')
logging.info(counter) logging.info(emoticounter)
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment