-
Notifications
You must be signed in to change notification settings - Fork 0
/
model003.py
110 lines (88 loc) · 3.56 KB
/
model003.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
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
# Based on http://visal.cs.cityu.edu.hk/static/pubs/conf/cvpr14w-hmlpe.pdf
import os
import time
import tflearn
from tflearn import layers
from tflearn.data_utils import image_preloader as preloader
import numpy as np
from PIL import Image
# For model saving
MODEL_ID = 3
WEIGHTS_FILE = 'weights/model_{:03d}'.format(MODEL_ID)
ONLY_TEST = False
# Configs
IMAGE_WIDTH, IMAGE_HEIGHT, IMAGE_CHANNELS = 720, 405, 3
X, _ = preloader('images.txt', image_shape=(IMAGE_HEIGHT, IMAGE_WIDTH), mode='file', categorical_labels=False)
# paths, joints = process(1)
# X = list()
# for path in paths:
# X.append(np.array(Image.open('data/frames/' + path)))
# X = np.array(X, dtype=np.float)
# X = X.reshape((-1, IMAGE_WIDTH, IMAGE_HEIGHT, 3))
# X = X.reshape((-1, IMAGE_WIDTH, IMAGE_HEIGHT))
Y = np.load('joints.npy')
for i, jo in enumerate(Y):
img = Image.open(X.array[i])
width, height = img.size
scales = [ IMAGE_WIDTH / width, IMAGE_HEIGHT / height]
Y[i][::2] = jo[::2] * scales[0]
Y[i][1::2] = jo[1::2] * scales[1]
# for index, image in enumerate(X.array[0:100]):
# img = Image.open(image)
# img = img.resize((IMAGE_WIDTH, IMAGE_HEIGHT), Image.ANTIALIAS)
#
# for i in range(0, 12, 2):
# img.putpixel((int(Y[index][i]), int(Y[index][i+1])), (255, 0, 0))
#
# print(str(index) + 'criado')
# img.save('./teste/' + str(index) + '.png', 'PNG')
# _, joints = process()
# Y = list()
# for joint in joints:
# Y.append(joint.flatten())
# Y = np.array(Y)
# print(Y[0])
# Network
net = layers.input_data([None, IMAGE_WIDTH, IMAGE_HEIGHT, IMAGE_CHANNELS])
net = layers.conv_2d(net, 64, 10, padding='valid', activation='relu')
net = layers.max_pool_2d(net, 10)
net = layers.conv_2d(net, 32, 6, padding='valid', activation='relu')
net = layers.conv_2d(net, 32, 6, padding='valid', activation='relu')
net = layers.max_pool_2d(net, 4)
net = layers.conv_2d(net, 16, 4, padding='valid', activation='relu')
net = layers.conv_2d(net, 16, 4, padding='valid', activation='relu')
net = layers.max_pool_2d(net, 2)
# net = layers.normalization.l2_normalize(net, 0)
net = layers.flatten(net)
net = layers.fully_connected(net, 2048, activation='relu')
net = layers.dropout(net, 0.25)
net = layers.fully_connected(net, 1024, activation='relu')
net = layers.dropout(net, 0.25)
net = layers.fully_connected(net, 18, activation='relu')
net = layers.regression(net, loss='mean_square', optimizer='adam')
# Model
model = tflearn.DNN(net, tensorboard_verbose=1)
if os.path.exists(WEIGHTS_FILE+'.index'):
print('========== Carregado =========')
model.load(WEIGHTS_FILE)
if ONLY_TEST:
print("=== test ===")
print(np.array(model.predict([X[10]]), dtype=np.uint))
print(np.array([Y[10]], dtype=np.uint))
else:
model.fit(X, Y, 250, validation_set=0.1, # 10% as validation
show_metric=True, snapshot_step=200, batch_size=10,
snapshot_epoch=False, run_id=WEIGHTS_FILE+ '::' +str(int(time.time())))
model.save(WEIGHTS_FILE)
else:
model.fit(X, Y, 250, validation_set=0.1, # 10% as validation
show_metric=True, snapshot_step=200, batch_size=10,
snapshot_epoch=False, run_id=WEIGHTS_FILE+ '::' +str(int(time.time())))
model.save(WEIGHTS_FILE)
if ONLY_TEST == False:
from util.result_check import render
for ind in range(0, 10):
predict = np.array(model.predict([X[ind]]), dtype=np.uint)
# ground = np.array([Y[ind]])
render(X.array[ind], predict[0], str(ind) + "_predict", (IMAGE_WIDTH, IMAGE_HEIGHT))
# render(X.array[ind], ground[0], str(ind) + "_ground")