forked from rlamarques/face-emotion-detection
-
Notifications
You must be signed in to change notification settings - Fork 0
/
count-images.py
82 lines (63 loc) · 2.41 KB
/
count-images.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
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
import os
import cv2
import math
import numpy as np
from shutil import copy2
import matplotlib.pyplot as plt
# Global variables
img_path = '/home/user/class/mc906/p4/dataset/croped/'
annotation_file = '/home/user/class/mc906/p4/dataset/data/label/label.lst'
f = open(os.path.join(annotation_file), "r")
labels = f.readlines()
emotion_dict = {}
for line in labels:
split = line.split(' ')
emotion_dict[split[0]] = split[7][:-1]
print('Emotion Dict', len(emotion_dict))
names = []
for original in os.listdir(img_path):
if original.endswith(".jpg"):
no_crop = original[6:]
names.append(no_crop)
print('Names', len(names))
# Makes sure that the cropped images are present in the label
new_dict = {}
for name in emotion_dict:
if name in names:
new_dict[name] = emotion_dict[name]
print('New dict', len(new_dict))
emotion_count = [0,0,0,0,0,0,0]
for emotion in new_dict.values():
emotion_count[int(emotion)] += 1
print(emotion_count)
print(sum(x for x in emotion_count))
translation = ['angry', 'disgust', 'fear', 'happy', 'sad', 'surprise', 'neutral']
txt = open('../image-count.txt', 'w')
distribution = []
for i, count in enumerate(emotion_count):
training = math.floor(0.6*count)
validation = test = math.floor(0.2*count)
total = training + validation + test
training += count - total
row = [training, validation, test]
distribution.append(row)
print(translation[i], 'training:', training, 'validation:', validation, 'test:', test)
txt.write(str(translation[i] +' ' + str(training) + ' ' + str(validation) + ' ' + str(test) + '\n'))
txt.close()
training_path = '/home/user/class/mc906/p4/dataset/organized/training/'
validation_path = '/home/user/class/mc906/p4/dataset/organized/validation/'
test_path = '/home/user/class/mc906/p4/dataset/organized/test/'
for cropped in os.listdir(img_path):
if cropped.endswith('.jpg'):
fname = cropped[6:]
emotion = int(emotion_dict[fname])
if distribution[emotion][0] > 0:
copy2(cropped, os.path.join(training_path, translation[emotion]))
distribution[emotion][0] -= 1
elif distribution[emotion][1] > 0:
copy2(cropped, os.path.join(validation_path, translation[emotion]))
distribution[emotion][1] -= 1
else:
copy2(cropped, os.path.join(test_path, translation[emotion]))
distribution[emotion][2] -= 1
print(distribution)