Commit 666f8d03 authored by tihmels's avatar tihmels
Browse files

AffectNet Pipeline

parent 5af0c9a4
import argparse
import logging
import os
import shutil
import sys
from shutil import copyfile
import pandas as pd
......@@ -34,17 +32,19 @@ parser.add_argument('-a', '--arousal', action='store', default='-2', dest='arous
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'}
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:
counter[e] = 0
emoticounter[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')
......@@ -58,17 +58,15 @@ def generate_records(csvfile):
if not os.path.exists(dest):
os.makedirs(dest)
else:
counter[e] = len(os.listdir(dest))
emoticounter[e] = len(os.listdir(dest))
data = zip(labels, imagepaths, valence, arousal)
for d in data:
if all(count >= arguments.limit for count in counter.values()):
break
emotion = emotidict.get(d[0])
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:
if emotion not in arguments.emotions or d[2] < arguments.valence or d[3] < arguments.arousal or emoticounter[
emotion] >= arguments.limit:
continue
img = d[1]
......@@ -78,18 +76,20 @@ def generate_records(csvfile):
logging.info(img + ' not found')
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):
copyfile(pathfrom, pathto)
logging.info('wrote file ' + pathto)
counter[emotion] += 1
if all(count >= arguments.limit for count in emoticounter.values()):
break
if __name__ == '__main__':
generate_records(csv)
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