Skip to content

Commit 022e8c2

Browse files
Laicheng0830Laicheng0830
committed
add pose estimation
1 parent 2d651fe commit 022e8c2

5 files changed

Lines changed: 394 additions & 1 deletion

File tree

‎tensorlayer/app/__init__.py‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,4 +2,5 @@
22
# -*- coding: utf-8 -*-
33

44
from .computer_vision_object_detection import *
5+
from .human_pose_estimation import *
56
from .computer_vision import *

‎tensorlayer/app/computer_vision.py‎

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22
# -*- coding: utf-8 -*-
33

44
from tensorlayer.app import YOLOv4, get_anchors, decode, filter_boxes
5+
from tensorlayer.app import CGCNN
56
import numpy as np
67
import tensorflow as tf
78
from tensorlayer import logging
@@ -41,6 +42,8 @@ def __init__(self, model_name='yolo4-mscoco'):
4142
self.model_name = model_name
4243
if self.model_name == 'yolo4-mscoco':
4344
self.model = YOLOv4(NUM_CLASS=80, pretrained=True)
45+
elif self.model_name == 'lcn':
46+
self.model = CGCNN(pretrained=True)
4447
else:
4548
raise ("The model does not support.")
4649

@@ -49,6 +52,8 @@ def __call__(self, input_data):
4952
batch_data = yolo4_input_processing(input_data)
5053
feature_maps = self.model(batch_data, is_train=False)
5154
output = yolo4_output_processing(feature_maps)
55+
elif self.model_name == 'lcn':
56+
output = self.model(input_data)
5257
else:
5358
raise NotImplementedError
5459

@@ -61,7 +66,7 @@ def __repr__(self):
6166

6267
@property
6368
def list(self):
64-
logging.info("The model name list: yolov4-mscoco")
69+
logging.info("The model name list: 'yolov4-mscoco', 'lcn'")
6570

6671

6772
def yolo4_input_processing(original_image):

‎tensorlayer/app/human_pose_estimation/LCN.py‎

Lines changed: 332 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,332 @@
1+
#! /usr/bin/python
2+
# -*- coding: utf-8 -*-
3+
""" LCN to estimate 3D human poses from 2D poses.
4+
5+
# Reference:
6+
- [pose_lcn](
7+
https://github.com/rujiewu/pose_lcn)
8+
9+
"""
10+
11+
import numpy as np
12+
import tensorflow as tf
13+
from tensorlayer.layers import Layer, Dropout, Dense, Input, BatchNorm, Reshape, Elementwise
14+
from tensorlayer.models import Model
15+
from tensorlayer import logging
16+
from .common import mask_weight, neighbour_matrix
17+
18+
BATCH_SIZE = 200
19+
M_0 = 17
20+
IN_F = 2
21+
22+
IN_JOINTS = 17
23+
OUT_JOINTS = 17
24+
F = 64
25+
NUM_LAYERS = 3
26+
weights_url = {'link': 'https://pan.baidu.com/s/1HBHWsAfyAlNaavw0iyUmUQ', 'password': 'ec07'}
27+
28+
29+
class Base_layer(Layer):
30+
31+
def __init__(
32+
self, F=F, in_joints=IN_JOINTS, out_joints=OUT_JOINTS, regularization=0.0, max_norm=True, residual=True,
33+
mask_type='locally_connected', neighbour_matrix=neighbour_matrix, init_type='ones', in_F=IN_F
34+
):
35+
super().__init__()
36+
self.F = F
37+
self.in_joints = in_joints
38+
self.regularizers = []
39+
self.regularization = regularization
40+
self.max_norm = max_norm
41+
self.out_joints = out_joints
42+
self.residual = residual
43+
self.mask_type = mask_type
44+
45+
self.init_type = init_type
46+
self.in_F = in_F
47+
48+
assert neighbour_matrix.shape[0] == neighbour_matrix.shape[1]
49+
assert neighbour_matrix.shape[0] == in_joints
50+
self.neighbour_matrix = neighbour_matrix
51+
52+
self._initialize_mask()
53+
54+
def _initialize_mask(self):
55+
"""
56+
Parameter
57+
mask_type
58+
locally_connected
59+
locally_connected_learnable
60+
init_type
61+
same: use L to init learnable part in mask
62+
ones: use 1 to init learnable part in mask
63+
random: use random to init learnable part in mask
64+
"""
65+
if 'locally_connected' in self.mask_type:
66+
assert self.neighbour_matrix is not None
67+
L = self.neighbour_matrix.T
68+
assert L.shape == (self.in_joints, self.in_joints)
69+
if 'learnable' not in self.mask_type:
70+
self.mask = tf.constant(L)
71+
else:
72+
if self.init_type == 'same':
73+
initializer = L
74+
elif self.init_type == 'ones':
75+
initializer = tf.initializers.ones
76+
elif self.init_type == 'random':
77+
initializer = tf.random.uniform
78+
var_mask = tf.Variable(
79+
name='mask', shape=[self.in_joints, self.out_joints] if self.init_type != 'same' else None,
80+
dtype=tf.float32, initial_value=initializer
81+
)
82+
var_mask = tf.nn.softmax(var_mask, axis=0)
83+
self.mask = var_mask * tf.constant(L != 0, dtype=tf.float32)
84+
85+
def _get_weights(self, name, initializer, shape, regularization=True, trainable=True):
86+
var = tf.Variable(initial_value=initializer(shape=shape, dtype=tf.float32), name=name, trainable=True)
87+
if regularization:
88+
self.regularizers.append(tf.nn.l2_loss(var))
89+
if trainable is True:
90+
if self._trainable_weights is None:
91+
self._trainable_weights = list()
92+
self._trainable_weights.append(var)
93+
else:
94+
if self._nontrainable_weights is None:
95+
self._nontrainable_weights = list()
96+
self._nontrainable_weights.append(var)
97+
return var
98+
99+
def kaiming(self, shape, dtype):
100+
"""Kaiming initialization as described in https://arxiv.org/pdf/1502.01852.pdf
101+
102+
Args
103+
shape: dimensions of the tf array to initialize
104+
dtype: data type of the array
105+
partition_info: (Optional) info about how the variable is partitioned.
106+
See https://github.com/tensorflow/tensorflow/blob/master/tensorflow/python/ops/init_ops.py#L26
107+
Needed to be used as an initializer.
108+
Returns
109+
Tensorflow array with initial weights
110+
"""
111+
return (tf.random.truncated_normal(shape, dtype=dtype) * tf.sqrt(2 / float(shape[0])))
112+
113+
def mask_weights(self, weights):
114+
return mask_weight(weights)
115+
116+
117+
class Mask_layer(Base_layer):
118+
119+
def __init__(self, in_channels=17, out_channels=None, name=None):
120+
super().__init__()
121+
self.in_channels = in_channels
122+
self.out_channels = out_channels
123+
self.w_name, self.b_name = name
124+
125+
if self.in_channels:
126+
self.build(None)
127+
self._built = True
128+
129+
def build(self, inputs_shape):
130+
if self.in_channels is None:
131+
self.in_channels = inputs_shape[1]
132+
133+
self.weight = self._get_weights(
134+
self.w_name, self.kaiming, [self.in_channels, self.out_channels], regularization=self.regularization != 0
135+
)
136+
self.bias = self._get_weights(
137+
self.b_name, self.kaiming, [self.out_channels], regularization=self.regularization != 0
138+
) # equal to b2leaky_relu
139+
self.weight = tf.clip_by_norm(self.weight, 1) if self.max_norm else self.weight
140+
141+
self.weight = self.mask_weights(self.weight)
142+
143+
def forward(self, x):
144+
outputs = tf.matmul(x, self.weight) + self.bias
145+
return outputs
146+
147+
148+
class End_layer(Base_layer):
149+
150+
def __init__(self):
151+
super().__init__()
152+
153+
def build(self, inputs_shape):
154+
pass
155+
156+
def forward(self, inputs):
157+
x, y = inputs
158+
x = tf.reshape(x, [-1, self.in_joints, self.in_F]) # [N, J, 3]
159+
y = tf.reshape(y, [-1, self.out_joints, 3]) # [N, J, 3]
160+
y = tf.concat([x[:, :, :2] + y[:, :, :2], tf.expand_dims(y[:, :, 2], axis=-1)], axis=2) # [N, J, 3]
161+
y = tf.reshape(y, [-1, self.out_joints * 3])
162+
return y
163+
164+
165+
def batch_normalization_warp(y):
166+
_, output_size = y.get_shape()
167+
output_size = int(output_size)
168+
out_F = int(output_size / IN_JOINTS)
169+
170+
y = Reshape([-1, IN_JOINTS, out_F])(y)
171+
y = BatchNorm(act='lrelu')(y)
172+
y = Reshape([-1, output_size])(y)
173+
return y
174+
175+
176+
def two_linear_train(inputs, idx):
177+
"""
178+
Make a bi-linear block with optional residual connection
179+
180+
Args
181+
xin: the batch that enters the block
182+
idx: integer. Number of layer (for naming/scoping)
183+
Returns
184+
y: the batch after it leaves the block
185+
"""
186+
187+
output_size = IN_JOINTS * F
188+
189+
# Linear 1
190+
input_size1 = int(inputs.get_shape()[1])
191+
output = Mask_layer(in_channels=input_size1, out_channels=output_size, name=["w2" + str(idx),
192+
"b2" + str(idx)])(inputs)
193+
output = batch_normalization_warp(output)
194+
output = Dropout(keep=0.8)(output)
195+
196+
# Linear 2
197+
input_size2 = int(output.get_shape()[1])
198+
output = Mask_layer(in_channels=input_size2, out_channels=output_size, name=["w3_" + str(idx),
199+
"b3_" + str(idx)])(output)
200+
output = batch_normalization_warp(output)
201+
output = Dropout(keep=0.8)(output)
202+
203+
# Residual every 2 blocks
204+
output = Elementwise(combine_fn=tf.add)([inputs, output])
205+
206+
return output
207+
208+
209+
def cgcnn_train():
210+
input_layer = Input(shape=(BATCH_SIZE, M_0 * IN_F))
211+
212+
# === First layer===
213+
output = Mask_layer(in_channels=IN_JOINTS * IN_F, out_channels=IN_JOINTS * F, name=["w1", "b1"])(input_layer)
214+
215+
output = batch_normalization_warp(output)
216+
output = Dropout(keep=0.8)(output)
217+
218+
# === Create multiple bi-linear layers ===
219+
for idx in range(NUM_LAYERS):
220+
output = two_linear_train(output, idx)
221+
222+
# === Last layer ===
223+
input_size4 = int(output.get_shape()[1])
224+
output = Mask_layer(in_channels=input_size4, out_channels=OUT_JOINTS*3, name=["w4", "b4"])(output)
225+
226+
# === End linear model ===
227+
output = End_layer()([input_layer,output])
228+
229+
network = Model(inputs=input_layer, outputs=output)
230+
231+
return network
232+
233+
234+
# inference
235+
def two_linear_inference(xin):
236+
"""
237+
Make a bi-linear block with optional residual connection
238+
239+
Args
240+
xin: the batch that enters the block
241+
y: the batch after it leaves the block
242+
"""
243+
244+
output_size = IN_JOINTS * F
245+
246+
# Linear 1
247+
output = Dense(n_units=output_size, act=None)(xin)
248+
output = batch_normalization_warp(output)
249+
# output = Dropout(keep=0.8)(output)
250+
251+
# Linear 2
252+
output = Dense(n_units=output_size, act=None)(output)
253+
output = batch_normalization_warp(output)
254+
# output = Dropout(keep=0.8)(output)
255+
256+
# Residual every 2 blocks
257+
y = Elementwise(tf.add)([xin, output])
258+
259+
return y
260+
261+
262+
def cgcnn_inference():
263+
input_layer = Input(shape=(BATCH_SIZE, M_0 * IN_F))
264+
265+
# === First layer===
266+
output = Dense(n_units=IN_JOINTS * F, act=None)(input_layer)
267+
output = batch_normalization_warp(output)
268+
# output = Dropout(keep=0.8)(output)
269+
270+
# === Create multiple bi-linear layers ===
271+
for i in range(3):
272+
output = two_linear_inference(output)
273+
274+
# === Last layer ===
275+
output = Dense(n_units=OUT_JOINTS * 3, act=None)(output)
276+
277+
output = End_layer()([input_layer, output])
278+
279+
network = Model(inputs=input_layer, outputs=output)
280+
return network
281+
282+
283+
def restore_params(network, model_path='model.npz'):
284+
logging.info("Restore pre-trained weights")
285+
286+
try:
287+
npz = np.load(model_path, allow_pickle=True)
288+
except:
289+
print("Download the model file, placed in the /model ")
290+
print("Weights download: ", weights_url['link'], "password:", weights_url['password'])
291+
292+
txt_path = 'model/pose_config.txt'
293+
f = open(txt_path, "r")
294+
line = f.readlines()
295+
for i in range(len(line)):
296+
# mask weights
297+
if len(npz[line[i].strip()].shape) == 2:
298+
_weight = mask_weight(npz[line[i].strip()])
299+
else:
300+
_weight = npz[line[i].strip()]
301+
network.all_weights[i].assign(_weight)
302+
logging.info(" Loading weights %s in %s" % (network.all_weights[i].shape, network.all_weights[i].name))
303+
304+
305+
def CGCNN(pretrained=True):
306+
"""Pre-trained LCN model.
307+
308+
Parameters
309+
------------
310+
pretrained : boolean
311+
Whether to load pretrained weights. Default False.
312+
313+
Examples
314+
---------
315+
Object Detection with YOLOv4, see `computer_vision.py
316+
<https://github.com/tensorlayer/tensorlayer/blob/master/tensorlayer/app/computer_vision.py>`__
317+
With TensorLayer
318+
319+
>>> # get the whole model, without pre-trained LCN parameters
320+
>>> lcn = tl.app.CGCNN(pretrained=False)
321+
>>> # get the whole model, restore pre-trained LCN parameters
322+
>>> lcn = tl.app.CGCNN(pretrained=True)
323+
>>> # use for inferencing
324+
>>> output = lcn(img, is_train=False)
325+
326+
"""
327+
if pretrained:
328+
network = cgcnn_inference()
329+
restore_params(network, model_path='model/model.npz')
330+
else:
331+
network = cgcnn_train()
332+
return network

‎tensorlayer/app/human_pose_estimation/__init__.py‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,5 @@
1+
#! /usr/bin/python
2+
# -*- coding: utf-8 -*-
3+
4+
from .common import *
5+
from .LCN import CGCNN

0 commit comments

Comments
 (0)