-
Notifications
You must be signed in to change notification settings - Fork 2
/
dataset_pickler.py
63 lines (38 loc) · 1.25 KB
/
dataset_pickler.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
#!/usr/bin/env python
# coding: utf-8
# In[1]:
import cv2
import os
from sklearn.model_selection import train_test_split
import numpy as np
# In[2]:
dataset_path = './dataset'
image_size_y = 50
image_size_x = 80
image_data = []
image_class = []
# In[3]:
for i, dir in enumerate(os.listdir(dataset_path)):
class_dir = (dataset_path + '/' + dir)
for image in os.listdir(class_dir):
image_file = class_dir + '/' + image
img = cv2.imread(image_file)
res = cv2.resize(img, dsize=(image_size_y, image_size_x), interpolation=cv2.INTER_CUBIC)
image_data.append(res)
image_class.append((i - 1) * -1)
image_data = np.array(image_data, dtype='float32')
image_class = np.array(image_class, dtype='float64')
# In[4]:
train_image, test_image, train_label, test_label = train_test_split(image_data,
image_class,
test_size=0.2)
# In[5]:
import pickle
file = open('test_data.pkl', 'wb')
test_set = (test_image, test_label)
pickle.dump(test_set, file)
file.close()
file = open('train_data.pkl', 'wb')
train_set = (train_image, train_label)
pickle.dump(train_set, file)
file.close()