-
Notifications
You must be signed in to change notification settings - Fork 1
/
process_emotions.py
49 lines (32 loc) · 1.23 KB
/
process_emotions.py
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
import os
import sys
import cv2
import numpy as np
import pandas as pd
from tqdm import tqdm
from EmoPy.src.fermodel import FERModel
if __name__ == "__main__":
if len(sys.argv) < 3:
sys.exit("Usage: python3 process_emotions.py <video directory> <results directory>")
video_dir = sys.argv[1]
result_dir = sys.argv[2]
files = [(os.path.join(video_dir,f), os.path.join(result_dir,os.path.splitext(f)[0]+'.csv')) for f in os.listdir(video_dir) if os.path.isfile(os.path.join(video_dir,f))]
target_emotions = ['anger', 'fear', 'calm', 'sadness', 'happiness', 'surprise', 'disgust']
model = FERModel(target_emotions, verbose=True)
temp_file = os.path.join(result_dir,'temp_img.jpg')
for (inF, outF) in files:
cap = cv2.VideoCapture(inF)
data = []
flag = True
frameId = 0
while cap.isOpened() and flag:
flag, frame = cap.read()
if flag:
cv2.imwrite(temp_file,frame)
v = model.predict(temp_file)
data.append([frameId,v])
print(frameId)
frameId += 1
cap.release()
df = pd.DataFrame(data, columns=['frameId','emotion'])
df.to_csv(outF, index=False)