tf-model-2/create_model.py at master · Coloriser/tf-model-2 · 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
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
from __future__ import division, print_function, absolute_import

import tflearn
from tflearn.layers.core import input_data, dropout, fully_connected
from tflearn.layers.conv import conv_2d, max_pool_2d
from tflearn.layers.normalization import local_response_normalization
from tflearn.layers.estimator import regression

import numpy as np
import argparse

import sys
sys.path.insert(0, './helper_modules')
import helper_functions as hf

def parse_arguments():						#argument parser -d for the pathlist
    parser = argparse.ArgumentParser(description='Trains the model, to be used after running pre-works')
    parser.add_argument('-a', help='to train based on a-channel', required=False,action="store_true", default=False)
    parser.add_argument('-b', help='to train based on b-channel', required=False,action="store_true", default=False)
    args = parser.parse_args()
    return args


def make_model(x, y):

	print("X :",x.shape)
	print("Y :",y.shape)


	# Building convolutional network
	network = input_data(shape=[None, x.shape[1], x.shape[2], 1], name='input')

	#1
	network = fully_connected(network, 128, activation='sigmoid')
	network = dropout(network, 0.8)
	print(network)
	#2
	network = fully_connected(network, 128, activation='sigmoid')
	network = dropout(network, 0.8)
	print(network)
	#3
	network = fully_connected(network, 128, activation='sigmoid')
	network = dropout(network, 0.8)
	print(network)
	#4
	network = fully_connected(network, 128, activation='sigmoid')
	network = dropout(network, 0.8)
	#5
	network = fully_connected(network, 128, activation='sigmoid')
	network = dropout(network, 0.8)
	#6
	network = fully_connected(network, 128, activation='sigmoid')
	network = dropout(network, 0.8)
	#7
	network = fully_connected(network, 128, activation='sigmoid')
	network = dropout(network, 0.8)
	#8
	network = fully_connected(network, 128, activation='sigmoid')
	network = dropout(network, 0.8)
	#9
	network = fully_connected(network, 128, activation='sigmoid')
	network = dropout(network, 0.8)
	#10
	network = fully_connected(network, 128, activation='sigmoid')
	network = dropout(network, 0.8)

	network = fully_connected(network, y.shape[1], activation='sigmoid')
	network = regression(network, optimizer='adam', learning_rate=0.01,
	                     loss='categorical_crossentropy', name='target')

	# Training
	model = tflearn.DNN(network, tensorboard_verbose=2)
	model.fit({'input': x}, {'target': y} , n_epoch=100)

	return model


def prereq_load_and_compute( mode , SIFT=False):
	if SIFT==True:
		print("SIFT")
		paths = hf.load_sift_paths('train')
	else:
		print("BRISK")
		paths = hf.load_brisk_paths('train')
	print("loading features...")
	features = hf.load_features(paths)
	print(str(len(features)) + " items loaded.")
	print("Normalizing features")
	modified_feature_arr = hf.normalize_array(features)
	No_Of_Test_Items = len(modified_feature_arr)

	if mode=='a':
		a_channel_paths = hf.load_a_channel_chroma_paths('train')
		print("loading a channel chroma...")
		a_channel_chromas = hf.load_a_channel_chroma(a_channel_paths)
		print(str(len(a_channel_chromas)) + " items loaded.")
		train_y_channel = np.array(a_channel_chromas).reshape(No_Of_Test_Items,-1)

	else:
		b_channel_paths = hf.load_b_channel_chroma_paths('train')
		print("loading b channel chroma...")
		b_channel_chromas = hf.load_b_channel_chroma(b_channel_paths)
		print(str(len(b_channel_chromas)) + " items loaded.")
		train_y_channel = np.array(b_channel_chromas).reshape(No_Of_Test_Items, -1)

	train_y_channel = train_y_channel+128
	train_y_channel = train_y_channel/256.0

	print("modifying the shape of input and output")
	train_x = np.array(modified_feature_arr).reshape([No_Of_Test_Items, modified_feature_arr[0].shape[0], modified_feature_arr[0].shape[1], 1])

	print("Pickling shapes")
	hf.pickle_shape(train_x,train_y_channel)

	print("train_x shape: ",train_x.shape)
	print("train_y shape: ",train_y_channel.shape)

	return train_x, train_y_channel

def make_a_model(  ):

	train_x, train_y_a_channel = prereq_load_and_compute( mode='a' , SIFT=True)

	print("Generating A channel model")
	model_a_channel = make_model(train_x, train_y_a_channel)
	model_a_channel.save("model/a_channel.model")

def make_b_model():

	train_x, train_y_b_channel = prereq_load_and_compute( mode='b' , SIFT=True)

	print("Generating B channel model")
	model_b_channel = make_model(train_x, train_y_b_channel)
	model_b_channel.save("model/b_channel.model")


def main():
	args = parse_arguments()
	# print(args)
	if args.a:
		print("Training model based on a-channel")
		make_a_model()
	if args.b:
		print("Training model based on b-channel")
		make_b_model()
	if not args.a and not args.b:
		print("ERROR: use -h for HELP")

main()