-
Notifications
You must be signed in to change notification settings - Fork 0
/
prepare_data.py
46 lines (39 loc) · 943 Bytes
/
prepare_data.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
import tensorflow as tf
import numpy as np
import random
mnist = tf.keras.datasets.mnist
(x_train, y_train), (x_test, y_test) = mnist.load_data()
def return_data(numbers):
Xtrain = []
Ytrain = []
for i in range(len(x_train)):
Xtrain.append(x_train[i].reshape(1,28,28))
temp = []
for j in range(10):
if y_train[i] == j:
temp.append(1)
else:
temp.append(0)
Ytrain.append(temp)
TrainIndices = []
for i in range(len(y_train)):
label = np.argmax(Ytrain[i])
if label in numbers:
TrainIndices.append(i)
Xtest = []
Ytest = []
for i in range(len(x_test)):
Xtest.append(x_test[i].reshape(1,28,28))
temp = []
for j in range(10):
if y_test[i] == j:
temp.append(1)
else:
temp.append(0)
Ytest.append(temp)
TestIndices = []
for i in range(len(y_test)):
label = np.argmax(Ytest[i])
if label in numbers:
TestIndices.append(i)
return Xtrain,Ytrain,Xtest,Ytest,TrainIndices,TestIndices