Skip to content
This repository was archived by the owner on Oct 31, 2023. It is now read-only.

Commit 723ab20

Browse files
JohnqczhangJohnqczhang
committed
refine body_uv_rcnn
1 parent 8a7f239 commit 723ab20

13 files changed

Lines changed: 575 additions & 592 deletions

‎detectron/core/config.py‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -876,6 +876,9 @@
876876
# Number of patches in the dataset
877877
__C.BODY_UV_RCNN.NUM_PATCHES = -1
878878

879+
# Number of semantic parts used to sample annotation points
880+
__C.BODY_UV_RCNN.NUM_SEMANTIC_PARTS = 14
881+
879882
# Number of stacked Conv layers in body UV head
880883
__C.BODY_UV_RCNN.NUM_STACKED_CONVS = 8
881884
# Dimension of the hidden representation output by the body UV head

‎detectron/core/test.py‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -948,16 +948,16 @@ def im_detect_body_uv(model, im_scale, boxes):
948948
# Removed squeeze calls due to singleton dimension issues
949949
CurAnnIndex = np.argmax(CurAnnIndex, axis=0)
950950
CurIndex_UV = np.argmax(CurIndex_UV, axis=0)
951-
CurIndex_UV = CurIndex_UV * (CurAnnIndex>0).astype(np.float32)
951+
CurIndex_UV = CurIndex_UV * (CurAnnIndex > 0).astype(np.float32)
952952

953953
output = np.zeros([3, int(by), int(bx)], dtype=np.float32)
954954
output[0] = CurIndex_UV
955955

956956
for part_id in range(1, K):
957957
CurrentU = CurU_uv[part_id]
958958
CurrentV = CurV_uv[part_id]
959-
output[1, CurIndex_UV==part_id] = CurrentU[CurIndex_UV==part_id]
960-
output[2, CurIndex_UV==part_id] = CurrentV[CurIndex_UV==part_id]
959+
output[1, CurIndex_UV == part_id] = CurrentU[CurIndex_UV == part_id]
960+
output[2, CurIndex_UV == part_id] = CurrentV[CurIndex_UV == part_id]
961961
outputs.append(output)
962962

963963
num_classes = cfg.MODEL.NUM_CLASSES

‎detectron/datasets/json_dataset.py‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -154,7 +154,7 @@ def _prep_roidb_entry(self, entry):
154154
(0, 3, self.num_keypoints), dtype=np.int32
155155
)
156156
if cfg.MODEL.BODY_UV_ON:
157-
entry['ignore_UV_body'] = np.empty((0), dtype=np.bool)
157+
entry['ignore_UV_body'] = np.empty((0), dtype=np.bool)
158158
# entry['Box_image_links_body'] = []
159159
# Remove unwanted fields that come from the json file (if they exist)
160160
for k in ['date_captured', 'url', 'license', 'file_name']:
@@ -200,7 +200,7 @@ def _add_gt_annotations(self, entry):
200200
valid_objs.append(obj)
201201
valid_segms.append(obj['segmentation'])
202202
###
203-
if 'dp_x' in obj.keys():
203+
if 'dp_x' in obj:
204204
valid_dp_x.append(obj['dp_x'])
205205
valid_dp_y.append(obj['dp_y'])
206206
valid_dp_I.append(obj['dp_I'])
@@ -216,7 +216,7 @@ def _add_gt_annotations(self, entry):
216216
valid_dp_masks.append([])
217217
###
218218
num_valid_objs = len(valid_objs)
219-
##
219+
220220
boxes = np.zeros((num_valid_objs, 4), dtype=entry['boxes'].dtype)
221221
gt_classes = np.zeros((num_valid_objs), dtype=entry['gt_classes'].dtype)
222222
gt_overlaps = np.zeros(
@@ -234,7 +234,7 @@ def _add_gt_annotations(self, entry):
234234
dtype=entry['gt_keypoints'].dtype
235235
)
236236
if cfg.MODEL.BODY_UV_ON:
237-
ignore_UV_body = np.zeros((num_valid_objs))
237+
ignore_UV_body = np.zeros((num_valid_objs), dtype=entry['ignore_UV_body'].dtype)
238238
#Box_image_body = [None]*num_valid_objs
239239

240240
im_has_visible_keypoints = False

‎detectron/datasets/roidb.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -121,7 +121,7 @@ def is_valid(entry):
121121
if cfg.MODEL.BODY_UV_ON and cfg.BODY_UV_RCNN.BODY_UV_IMS:
122122
# Exclude images with no body uv
123123
valid = valid and entry['has_body_uv']
124-
return valid
124+
return valid
125125

126126
num = len(roidb)
127127
filtered_roidb = [entry for entry in roidb if is_valid(entry)]

‎detectron/modeling/body_uv_rcnn_heads.py‎

Lines changed: 128 additions & 132 deletions
Original file line numberDiff line numberDiff line change
@@ -3,150 +3,146 @@
33
#
44
# This source code is licensed under the license found in the
55
# LICENSE file in the root directory of this source tree.
6-
#
6+
##############################################################################
7+
8+
"""Various network "heads" for dense human pose estimation in DensePose.
9+
10+
The design is as follows:
11+
12+
... -> RoI ----\ /-> mask output -> cls loss
13+
-> RoIFeatureXform -> body UV head -> patch output -> cls loss
14+
... -> Feature / \-> UV output -> reg loss
15+
Map
16+
17+
The body UV head produces a feature representation of the RoI for the purpose
18+
of dense semantic mask prediction, body surface patch prediction and body UV
19+
coordinates regression. The body UV output module converts the feature
20+
representation into heatmaps for dense mask, patch index and UV coordinates.
21+
"""
722

823
from __future__ import absolute_import
924
from __future__ import division
1025
from __future__ import print_function
1126
from __future__ import unicode_literals
1227

1328
from caffe2.python import core
14-
1529
from detectron.core.config import cfg
16-
30+
from detectron.utils.c2 import const_fill
1731
import detectron.modeling.ResNet as ResNet
1832
import detectron.utils.blob as blob_utils
1933

2034
# ---------------------------------------------------------------------------- #
21-
# Body UV heads
35+
# Body UV outputs and losses
2236
# ---------------------------------------------------------------------------- #
2337

24-
def add_body_uv_outputs(model, blob_in, dim, pref=''):
25-
####
26-
model.ConvTranspose(blob_in, 'AnnIndex_lowres'+pref, dim, 15,cfg.BODY_UV_RCNN.DECONV_KERNEL, pad=int(cfg.BODY_UV_RCNN.DECONV_KERNEL / 2 - 1), stride=2, weight_init=(cfg.BODY_UV_RCNN.CONV_INIT, {'std': 0.001}), bias_init=('ConstantFill', {'value': 0.}))
27-
####
28-
model.ConvTranspose(blob_in, 'Index_UV_lowres'+pref, dim, cfg.BODY_UV_RCNN.NUM_PATCHES+1,cfg.BODY_UV_RCNN.DECONV_KERNEL, pad=int(cfg.BODY_UV_RCNN.DECONV_KERNEL / 2 - 1), stride=2, weight_init=(cfg.BODY_UV_RCNN.CONV_INIT, {'std': 0.001}), bias_init=('ConstantFill', {'value': 0.}))
29-
####
30-
model.ConvTranspose(
31-
blob_in, 'U_lowres'+pref, dim, (cfg.BODY_UV_RCNN.NUM_PATCHES+1),
32-
cfg.BODY_UV_RCNN.DECONV_KERNEL,
33-
pad=int(cfg.BODY_UV_RCNN.DECONV_KERNEL / 2 - 1),
34-
stride=2,
35-
weight_init=(cfg.BODY_UV_RCNN.CONV_INIT, {'std': 0.001}),
36-
bias_init=('ConstantFill', {'value': 0.}))
37-
#####
38-
model.ConvTranspose(
39-
blob_in, 'V_lowres'+pref, dim, cfg.BODY_UV_RCNN.NUM_PATCHES+1,
38+
def add_body_uv_outputs(model, blob_in, dim):
39+
"""Add DensePose body UV specific outputs: heatmaps of dense mask, patch index
40+
and patch-specific UV coordinates. All dense masks are mapped to labels in
41+
[0, ... S] for S semantically meaningful body parts.
42+
"""
43+
# Apply ConvTranspose to the feature representation; results in 2x upsampling
44+
for name in ['AnnIndex', 'Index_UV', 'U', 'V']:
45+
if name == 'AnnIndex':
46+
dim_out = cfg.BODY_UV_RCNN.NUM_SEMANTIC_PARTS + 1
47+
else:
48+
dim_out = cfg.BODY_UV_RCNN.NUM_PATCHES + 1
49+
model.ConvTranspose(
50+
blob_in,
51+
name + '_lowres',
52+
dim,
53+
dim_out,
4054
cfg.BODY_UV_RCNN.DECONV_KERNEL,
4155
pad=int(cfg.BODY_UV_RCNN.DECONV_KERNEL / 2 - 1),
4256
stride=2,
4357
weight_init=(cfg.BODY_UV_RCNN.CONV_INIT, {'std': 0.001}),
44-
bias_init=('ConstantFill', {'value': 0.}))
45-
####
46-
blob_Ann_Index = model.BilinearInterpolation('AnnIndex_lowres'+pref, 'AnnIndex'+pref, cfg.BODY_UV_RCNN.NUM_PATCHES+1 , cfg.BODY_UV_RCNN.NUM_PATCHES+1, cfg.BODY_UV_RCNN.UP_SCALE)
47-
blob_Index = model.BilinearInterpolation('Index_UV_lowres'+pref, 'Index_UV'+pref, cfg.BODY_UV_RCNN.NUM_PATCHES+1 , cfg.BODY_UV_RCNN.NUM_PATCHES+1, cfg.BODY_UV_RCNN.UP_SCALE)
48-
blob_U = model.BilinearInterpolation('U_lowres'+pref, 'U_estimated'+pref, cfg.BODY_UV_RCNN.NUM_PATCHES+1 , cfg.BODY_UV_RCNN.NUM_PATCHES+1, cfg.BODY_UV_RCNN.UP_SCALE)
49-
blob_V = model.BilinearInterpolation('V_lowres'+pref, 'V_estimated'+pref, cfg.BODY_UV_RCNN.NUM_PATCHES+1 , cfg.BODY_UV_RCNN.NUM_PATCHES+1, cfg.BODY_UV_RCNN.UP_SCALE)
50-
###
51-
return blob_U,blob_V,blob_Index,blob_Ann_Index
52-
53-
54-
def add_body_uv_losses(model, pref=''):
55-
56-
## Reshape for GT blobs.
57-
model.net.Reshape( ['body_uv_X_points'], ['X_points_reshaped'+pref, 'X_points_shape'+pref], shape=( -1 ,1 ) )
58-
model.net.Reshape( ['body_uv_Y_points'], ['Y_points_reshaped'+pref, 'Y_points_shape'+pref], shape=( -1 ,1 ) )
59-
model.net.Reshape( ['body_uv_I_points'], ['I_points_reshaped'+pref, 'I_points_shape'+pref], shape=( -1 ,1 ) )
60-
model.net.Reshape( ['body_uv_Ind_points'], ['Ind_points_reshaped'+pref, 'Ind_points_shape'+pref], shape=( -1 ,1 ) )
61-
## Concat Ind,x,y to get Coordinates blob.
62-
model.net.Concat( ['Ind_points_reshaped'+pref,'X_points_reshaped'+pref, \
63-
'Y_points_reshaped'+pref],['Coordinates'+pref,'Coordinate_Shapes'+pref ], axis = 1 )
64-
##
65-
### Now reshape UV blobs, such that they are 1x1x(196*NumSamples)xNUM_PATCHES
66-
## U blob to
67-
##
68-
model.net.Reshape(['body_uv_U_points'], \
69-
['U_points_reshaped'+pref, 'U_points_old_shape'+pref],\
70-
shape=(-1,cfg.BODY_UV_RCNN.NUM_PATCHES+1,196))
71-
model.net.Transpose(['U_points_reshaped'+pref] ,['U_points_reshaped_transpose'+pref],axes=(0,2,1) )
72-
model.net.Reshape(['U_points_reshaped_transpose'+pref], \
73-
['U_points'+pref, 'U_points_old_shape2'+pref], \
74-
shape=(1,1,-1,cfg.BODY_UV_RCNN.NUM_PATCHES+1))
75-
## V blob
76-
##
77-
model.net.Reshape(['body_uv_V_points'], \
78-
['V_points_reshaped'+pref, 'V_points_old_shape'+pref],\
79-
shape=(-1,cfg.BODY_UV_RCNN.NUM_PATCHES+1,196))
80-
model.net.Transpose(['V_points_reshaped'+pref] ,['V_points_reshaped_transpose'+pref],axes=(0,2,1) )
81-
model.net.Reshape(['V_points_reshaped_transpose'+pref], \
82-
['V_points'+pref, 'V_points_old_shape2'+pref], \
83-
shape=(1,1,-1,cfg.BODY_UV_RCNN.NUM_PATCHES+1))
84-
###
85-
## UV weights blob
86-
##
87-
model.net.Reshape(['body_uv_point_weights'], \
88-
['Uv_point_weights_reshaped'+pref, 'Uv_point_weights_old_shape'+pref],\
89-
shape=(-1,cfg.BODY_UV_RCNN.NUM_PATCHES+1,196))
90-
model.net.Transpose(['Uv_point_weights_reshaped'+pref] ,['Uv_point_weights_reshaped_transpose'+pref],axes=(0,2,1) )
91-
model.net.Reshape(['Uv_point_weights_reshaped_transpose'+pref], \
92-
['Uv_point_weights'+pref, 'Uv_point_weights_old_shape2'+pref], \
93-
shape=(1,1,-1,cfg.BODY_UV_RCNN.NUM_PATCHES+1))
94-
95-
#####################
96-
### Pool IUV for points via bilinear interpolation.
97-
model.PoolPointsInterp(['U_estimated','Coordinates'+pref], ['interp_U'+pref])
98-
model.PoolPointsInterp(['V_estimated','Coordinates'+pref], ['interp_V'+pref])
99-
model.PoolPointsInterp(['Index_UV'+pref,'Coordinates'+pref], ['interp_Index_UV'+pref])
100-
101-
## Reshape interpolated UV coordinates to apply the loss.
102-
103-
model.net.Reshape(['interp_U'+pref], \
104-
['interp_U_reshaped'+pref, 'interp_U_shape'+pref],\
105-
shape=(1, 1, -1 , cfg.BODY_UV_RCNN.NUM_PATCHES+1))
106-
107-
model.net.Reshape(['interp_V'+pref], \
108-
['interp_V_reshaped'+pref, 'interp_V_shape'+pref],\
109-
shape=(1, 1, -1 , cfg.BODY_UV_RCNN.NUM_PATCHES+1))
110-
###
111-
112-
### Do the actual labels here !!!!
113-
model.net.Reshape( ['body_uv_ann_labels'], \
114-
['body_uv_ann_labels_reshaped' +pref, 'body_uv_ann_labels_old_shape'+pref], \
115-
shape=(-1, cfg.BODY_UV_RCNN.HEATMAP_SIZE , cfg.BODY_UV_RCNN.HEATMAP_SIZE))
116-
117-
model.net.Reshape( ['body_uv_ann_weights'], \
118-
['body_uv_ann_weights_reshaped' +pref, 'body_uv_ann_weights_old_shape'+pref], \
119-
shape=( -1 , cfg.BODY_UV_RCNN.HEATMAP_SIZE , cfg.BODY_UV_RCNN.HEATMAP_SIZE))
120-
###
121-
model.net.Cast( ['I_points_reshaped'+pref], ['I_points_reshaped_int'+pref], to=core.DataType.INT32)
122-
### Now add the actual losses
123-
## The mask segmentation loss (dense)
124-
probs_seg_AnnIndex, loss_seg_AnnIndex = model.net.SpatialSoftmaxWithLoss( \
125-
['AnnIndex'+pref, 'body_uv_ann_labels_reshaped'+pref,'body_uv_ann_weights_reshaped'+pref],\
126-
['probs_seg_AnnIndex'+pref,'loss_seg_AnnIndex'+pref], \
127-
scale=cfg.BODY_UV_RCNN.INDEX_WEIGHTS / cfg.NUM_GPUS)
128-
## Point Patch Index Loss.
129-
probs_IndexUVPoints, loss_IndexUVPoints = model.net.SoftmaxWithLoss(\
130-
['interp_Index_UV'+pref,'I_points_reshaped_int'+pref],\
131-
['probs_IndexUVPoints'+pref,'loss_IndexUVPoints'+pref], \
132-
scale=cfg.BODY_UV_RCNN.PART_WEIGHTS / cfg.NUM_GPUS, spatial=0)
133-
## U and V point losses.
134-
loss_Upoints = model.net.SmoothL1Loss( \
135-
['interp_U_reshaped'+pref, 'U_points'+pref, \
136-
'Uv_point_weights'+pref, 'Uv_point_weights'+pref], \
137-
'loss_Upoints'+pref, \
138-
scale=cfg.BODY_UV_RCNN.POINT_REGRESSION_WEIGHTS / cfg.NUM_GPUS)
58+
bias_init=const_fill(0.0)
59+
)
60+
# Increase heatmap output size via bilinear upsampling
61+
blob_outputs = []
62+
for name in ['AnnIndex', 'Index_UV', 'U', 'V']:
63+
blob_outputs.append(
64+
model.BilinearInterpolation(
65+
name + '_lowres',
66+
name + '_estimated' if name in ['U', 'V'] else name,
67+
cfg.BODY_UV_RCNN.NUM_PATCHES + 1,
68+
cfg.BODY_UV_RCNN.NUM_PATCHES + 1,
69+
cfg.BODY_UV_RCNN.UP_SCALE
70+
)
71+
)
72+
73+
return blob_outputs
74+
75+
76+
def add_body_uv_losses(model):
77+
"""Add DensePose body UV specific losses."""
78+
# Pool estimated IUV points via bilinear interpolation.
79+
for name in ['U', 'V', 'Index_UV']:
80+
model.PoolPointsInterp(
81+
[
82+
name + '_estimated' if name in ['U', 'V'] else name,
83+
'body_uv_coords_xy'
84+
],
85+
['interp_' + name]
86+
)
13987

140-
loss_Vpoints = model.net.SmoothL1Loss( \
141-
['interp_V_reshaped'+pref, 'V_points'+pref, \
142-
'Uv_point_weights'+pref, 'Uv_point_weights'+pref], \
143-
'loss_Vpoints'+pref, scale=cfg.BODY_UV_RCNN.POINT_REGRESSION_WEIGHTS / cfg.NUM_GPUS)
144-
## Add the losses.
145-
loss_gradients = blob_utils.get_loss_gradients(model, \
146-
[ loss_Upoints, loss_Vpoints, loss_seg_AnnIndex, loss_IndexUVPoints])
147-
model.losses = list(set(model.losses + \
148-
['loss_Upoints'+pref , 'loss_Vpoints'+pref , \
149-
'loss_seg_AnnIndex'+pref ,'loss_IndexUVPoints'+pref]))
88+
# Compute spatial softmax normalized probabilities, after which
89+
# cross-entropy loss is computed for semantic parts classification.
90+
probs_AnnIndex, loss_AnnIndex = model.net.SpatialSoftmaxWithLoss(
91+
[
92+
'AnnIndex',
93+
'body_uv_parts', 'body_uv_parts_weights'
94+
],
95+
['probs_AnnIndex', 'loss_AnnIndex'],
96+
scale=cfg.BODY_UV_RCNN.INDEX_WEIGHTS / cfg.NUM_GPUS
97+
)
98+
# Softmax loss for surface patch classification.
99+
probs_I_points, loss_I_points = model.net.SoftmaxWithLoss(
100+
['interp_Index_UV', 'body_uv_I_points'],
101+
['probs_I_points', 'loss_I_points'],
102+
scale=cfg.BODY_UV_RCNN.PART_WEIGHTS / cfg.NUM_GPUS,
103+
spatial=0
104+
)
105+
## Smooth L1 loss for each patch-specific UV coordinates regression.
106+
# Reshape U,V blobs of both interpolated and ground-truth to compute
107+
# summarized (instead of averaged) SmoothL1Loss.
108+
loss_UV = list()
109+
model.net.Reshape(
110+
['body_uv_point_weights'],
111+
['UV_point_weights', 'body_uv_point_weights_shape'],
112+
shape=(1, -1, cfg.BODY_UV_RCNN.NUM_PATCHES + 1)
113+
)
114+
for name in ['U', 'V']:
115+
# Reshape U/V coordinates of both interpolated points and ground-truth
116+
# points from (#points, #patches) to (1, #points, #patches).
117+
model.net.Reshape(
118+
['body_uv_' + name + '_points'],
119+
[name + '_points', 'body_uv_' + name + '_points_shape'],
120+
shape=(1, -1, cfg.BODY_UV_RCNN.NUM_PATCHES + 1)
121+
)
122+
model.net.Reshape(
123+
['interp_' + name],
124+
['interp_' + name + '_reshaped', 'interp_' + name + 'shape'],
125+
shape=(1, -1, cfg.BODY_UV_RCNN.NUM_PATCHES + 1)
126+
)
127+
# Compute summarized SmoothL1Loss of all points.
128+
loss_UV.append(
129+
model.net.SmoothL1Loss(
130+
[
131+
'interp_' + name + '_reshaped', name + '_points',
132+
'UV_point_weights', 'UV_point_weights'
133+
],
134+
'loss_' + name + '_points',
135+
scale=cfg.BODY_UV_RCNN.POINT_REGRESSION_WEIGHTS / cfg.NUM_GPUS
136+
)
137+
)
138+
# Add all losses to compute gradients
139+
loss_gradients = blob_utils.get_loss_gradients(
140+
model, [loss_AnnIndex, loss_I_points] + loss_UV
141+
)
142+
# Update model training losses
143+
model.AddLosses(
144+
['loss_' + name for name in ['AnnIndex', 'I_points', 'U_points', 'V_points']]
145+
)
150146

151147
return loss_gradients
152148

@@ -155,17 +151,17 @@ def add_body_uv_losses(model, pref=''):
155151
# Body UV heads
156152
# ---------------------------------------------------------------------------- #
157153

158-
def add_ResNet_roi_conv5_head_for_bodyUV(
159-
model, blob_in, dim_in, spatial_scale
160-
):
154+
def add_ResNet_roi_conv5_head_for_bodyUV(model, blob_in, dim_in, spatial_scale):
161155
"""Add a ResNet "conv5" / "stage5" head for body UV prediction."""
162156
model.RoIFeatureTransform(
163-
blob_in, '_[body_uv]_pool5',
157+
blob_in,
158+
'_[body_uv]_pool5',
164159
blob_rois='body_uv_rois',
165160
method=cfg.BODY_UV_RCNN.ROI_XFORM_METHOD,
166161
resolution=cfg.BODY_UV_RCNN.ROI_XFORM_RESOLUTION,
167162
sampling_ratio=cfg.BODY_UV_RCNN.ROI_XFORM_SAMPLING_RATIO,
168-
spatial_scale=spatial_scale)
163+
spatial_scale=spatial_scale
164+
)
169165
# Using the prefix '_[body_uv]_' to 'res5' enables initializing the head's
170166
# parameters using pretrained 'res5' parameters if given (see
171167
# utils.net.initialize_from_weights_file)
@@ -184,7 +180,7 @@ def add_ResNet_roi_conv5_head_for_bodyUV(
184180

185181

186182
def add_roi_body_uv_head_v1convX(model, blob_in, dim_in, spatial_scale):
187-
"""v1convX design: X * (conv)."""
183+
"""Add a DensePose body UV head. v1convX design: X * (conv)."""
188184
hidden_dim = cfg.BODY_UV_RCNN.CONV_HEAD_DIM
189185
kernel_size = cfg.BODY_UV_RCNN.CONV_HEAD_KERNEL
190186
pad_size = kernel_size // 2
@@ -208,7 +204,7 @@ def add_roi_body_uv_head_v1convX(model, blob_in, dim_in, spatial_scale):
208204
stride=1,
209205
pad=pad_size,
210206
weight_init=(cfg.BODY_UV_RCNN.CONV_INIT, {'std': 0.01}),
211-
bias_init=('ConstantFill', {'value': 0.})
207+
bias_init=const_fill(0.0)
212208
)
213209
current = model.Relu(current, current)
214210
dim_in = hidden_dim

0 commit comments

Comments
 (0)