Commit 666f8d03 authored by tihmels's avatar tihmels

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