-
Notifications
You must be signed in to change notification settings - Fork 2
/
preprocess.py
55 lines (42 loc) · 1.5 KB
/
preprocess.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
import numpy as np
import tensorflow as tf
from os import walk
from os.path import join
from shutil import copyfile
from hyperparameters import BATCH_SIZE
from random import shuffle
import cv2
from hyperparameters import IMAGE_HEIGHT, IMAGE_WIDTH
gpu_available = tf.test.is_gpu_available()
print("GPU Available: ", gpu_available)
data_paths = []
def get_batch(batch_idx):
batch = np.asarray([img2lab(file) for file in data_paths[batch_idx * BATCH_SIZE: (batch_idx + 1) * BATCH_SIZE]])
return (batch[:, :, :, 0].reshape((-1, IMAGE_HEIGHT, IMAGE_WIDTH, 1)).astype('float32'),
batch[:, :, :, 1:].astype('float32'))
def fetch_data():
"""
walk through data set directory
:return: None, you should save all images to one directory
"""
global data_paths
flip = 1
data_dir = "SUN2012/Images/"
for root, subfolder, files in walk(data_dir):
for file in files:
if file.endswith('.jpg'):
if flip == 1:
data_paths.append(join(root, file))
flip = 0
# if len(data_paths) == 30:
# return 30
# copyfile(join(root, file), join("preprocessed/", file))
else:
flip = 1
return len(data_paths)
def shuffle_data():
shuffle(data_paths)
def img2lab(file_name):
img = cv2.imread(file_name)
lab_img = cv2.cvtColor(img, cv2.COLOR_BGR2LAB)
return ((lab_img / [100, 1, 1]) * [255, 1, 1]) - [0, 128, 128]