NeuralWeightVirtualization/GTSRB_data.py at master · learning1234embed/NeuralWeightVirtualization · GitHub

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
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
from __future__ import print_function
from PIL import Image
import numpy as np
import os

train_path = '...'
test_path = '...'

#"""
GTSRB_train_data = np.load('GTSRB_train_data.npy')
GTSRB_train_label = np.load('GTSRB_train_label.npy')

GTSRB_test_data = np.load('GTSRB_test_data.npy')
GTSRB_test_label = np.load('GTSRB_test_label.npy')
#"""

def train_set():
	return GTSRB_train_data, GTSRB_train_label

def test_set():
	return GTSRB_test_data, GTSRB_test_label

def validation_set():
	return None, None

def get_GTSRB_test_batch(dir_path):
	GTSRB_batch = []
	labels = []

	csv_filepath = os.path.join(dir_path, 'GT-final_test.csv')
	if not os.path.exists(csv_filepath):
        	return

	print(csv_filepath)

	with open(csv_filepath) as f:
		for line in f:
			lists = line.split(';')
			img_filepath = os.path.join(dir_path, lists[0])
			new_img_filepath = os.path.join(dir_path, os.path.basename(img_filepath))
			new_img_filepath = os.path.splitext(new_img_filepath)[0] + '.png'
			if not os.path.exists(new_img_filepath):
		                continue

			img = Image.open(new_img_filepath)
			GTSRB = np.array(img)
			GTSRB = GTSRB.reshape((np.size(GTSRB)))
			normalized_GTSRB = GTSRB / np.linalg.norm(GTSRB)
			label = int(lists[7])
			GTSRB_batch.append(normalized_GTSRB)
			labels.append(label)

	GTSRB_batch = np.stack((GTSRB_batch))
	print(GTSRB_batch.shape)

	labels_one_hot = np.zeros((GTSRB_batch.shape[0], np.max(labels)+1))
	labels_one_hot[np.arange(GTSRB_batch.shape[0]), labels] = 1
	print(labels_one_hot.shape)

	return GTSRB_batch, labels_one_hot

def get_GTSRB_train_batch(dir_path):
	GTSRB_batch = []
	labels = []
	label = 0

	for (dirpath, dirnames, filenames) in sorted(os.walk(dir_path)):
		if os.path.basename(dirpath) == '':
			continue

		for (dirpath2, dirnames2, filenames2) in sorted(os.walk(dirpath)):
			for file in sorted(filenames2):
				filepath = os.path.join(dirpath2, file)
				img = Image.open(filepath)
				GTSRB = np.array(img)
				GTSRB = GTSRB.reshape((np.size(GTSRB)))
				normalized_GTSRB = GTSRB / np.linalg.norm(GTSRB)
				GTSRB_batch.append(normalized_GTSRB)
				labels.append(label)

		label += 1

	GTSRB_batch = np.vstack((GTSRB_batch))
	labels_one_hot = np.zeros((GTSRB_batch.shape[0], label))
	labels_one_hot[np.arange(GTSRB_batch.shape[0]), labels] = 1

	print(labels_one_hot.shape)

	return GTSRB_batch, labels_one_hot

def create_data_files():
	GTSRB_train_data, GTSRB_train_label = get_GTSRB_train_batch(train_path)
	GTSRB_test_data, GTSRB_test_label = get_GTSRB_test_batch(test_path)

	np.save('GTSRB_train_data', GTSRB_train_data)
	np.save('GTSRB_train_label', GTSRB_train_label)
	np.save('GTSRB_test_data', GTSRB_test_data)
	np.save('GTSRB_test_label', GTSRB_test_label)

def main():
	#create_data_files()
	print(train_set()[0])
	print(train_set()[0].shape)
	print(train_set()[1])
	print(train_set()[1].shape)

if __name__ == '__main__':
        main()