def _add_roi_mask_head(model, add_roi_mask_head_func, blob_in, dim_in,
spatial_scale_in):
"""Add a mask prediction head to the model."""
# Capture model graph before adding the mask head
bbox_net = copy.deepcopy(model.net.Proto())
# Add the mask head
blob_mask_head, dim_mask_head = add_roi_mask_head_func(
model, blob_in, dim_in, spatial_scale_in)
# Add the mask output
blob_mask = mask_rcnn_heads.add_mask_rcnn_outputs(model, blob_mask_head,
dim_mask_head)
if not model.train: # == inference
# Inference uses a cascade of box predictions, then mask predictions.
# This requires separate nets for box and mask prediction.
# So we extract the mask prediction net, store it as its own network,
# then restore model.net to be the bbox-only network
model.mask_net, blob_mask = c2_utils.SuffixNet('mask_net', model.net,
len(bbox_net.op),
blob_mask)
model.net._net = bbox_net
loss_gradients = None
else:
loss_gradients = mask_rcnn_heads.add_mask_rcnn_losses(model, blob_mask)
return loss_gradients
def _add_roi_body_uv_head(
model, add_roi_body_uv_head_func, blob_in, dim_in, spatial_scale_in
):
"""Add a body UV prediction head to the model."""
# Capture model graph before adding the mask head
bbox_net = copy.deepcopy(model.net.Proto())
# Add the body UV head
blob_body_uv_head, dim_body_uv_head = add_roi_body_uv_head_func(
model, blob_in, dim_in, spatial_scale_in
)
# Add the body UV output
blobs_body_uv = body_uv_rcnn_heads.add_body_uv_outputs(
model, blob_body_uv_head, dim_body_uv_head
)
if not model.train: # == inference
# Inference uses a cascade of box predictions, then body uv predictions
# This requires separate nets for box and body uv prediction.
# So we extract the keypoint prediction net, store it as its own
# network, then restore model.net to be the bbox-only network
model.body_uv_net, body_uv_blob_out = c2_utils.SuffixNet(
'body_uv_net', model.net, len(bbox_net.op), blobs_body_uv
)
model.net._net = bbox_net
loss_gradients = None
else:
loss_gradients = body_uv_rcnn_heads.add_body_uv_losses(model)
return loss_gradients
def add_fpn_onto_conv_body(model,
conv_body_func,
fpn_level_info_func,
P2only=False):
"""Add the specified conv body to the model and then add FPN levels to it.
"""
# Note: blobs_conv is in revsersed order: [fpn5, fpn4, fpn3, fpn2]
# similarly for dims_conv: [2048, 1024, 512, 256]
# similarly for spatial_scales_fpn: [1/32, 1/16, 1/8, 1/4]
conv_body_func(model)
blobs_fpn, dim_fpn, spatial_scales_fpn = add_fpn(model,
fpn_level_info_func())
if cfg.MODEL.SIBLING_BACKBONE_ON:
# Resnet stage to fork backbone weights
fork_at = cfg.SIBLING.FORK_AT
assert fork_at in [0, 2, 3, 4, 5]
prefix_len = 0
# No shared weights
if fork_at == 0:
fork_node = core.ScopedName("data".format(fork_at, fork_at))
# Fork weights at stage `fork_at`
else:
fork_node = core.ScopedName("res{}_{}_sum".format(
fork_at, fork_at))
ops = model.net.Proto().op
while (fork_node not in ops[prefix_len].output):
prefix_len += 1
prefix_len += 2
temp_net, _ = c2_utils.SuffixNet('temp_net', model.net, prefix_len,
blobs_fpn)
# Preffix sibling backbone
temp_net_preffixed, _ = c2_utils.RenameNet("temp_net_preffixed",
temp_net,
cfg.SIBLING.PREFFIX,
excluded_nodes=[fork_node])
model.AddParams([
core.BlobReference(input_name)
for op in temp_net_preffixed.Proto().op for input_name in op.input
if input_name[-2] == "_"
])
# Merge the backbones
model.net = c2_utils.MergeNets("net", [model.net, temp_net_preffixed])
del temp_net
del temp_net_preffixed
if P2only:
# use only the finest level
return blobs_fpn[-1], dim_fpn, spatial_scales_fpn[-1]
else:
# use all levels
return blobs_fpn, dim_fpn, spatial_scales_fpn
def _add_roi_track_head(
model, add_roi_track_head_func, blob_in, dim_in, spatial_scale_in
):
"""Add a track prediction head to the model."""
# Capture model graph before adding the track head
bbox_net = copy.deepcopy(model.net.Proto())
# Add the track head
blob_track_head, dim_track_head = add_roi_track_head_func(
model, blob_in, dim_in, spatial_scale_in
)
# Add the track output
blob_track = track_rcnn_heads.add_track_outputs(
model, blob_track_head, dim_track_head
)
if not model.train: # == inference
# Inference uses a cascade of box predictions, then object association predictions
# This requires separate nets for box and track prediction.
# So we extract the track prediction net, store it as its own
# # network, then restore model.net to be the bbox-only network
if cfg.MODEL.SIBLING_BACKBONE_ON and 'track' in cfg.SIBLING.HEADS:
track_net_temp, _ = c2_utils.SuffixNet(
'track_net_temp', model.net, len(bbox_net.op), blob_track
)
model.track_net, _ = c2_utils.RenameNet(
"track_net", track_net_temp, cfg.SIBLING.PREFFIX, excluded_nodes=[core.ScopedName("track_rois_fpn{}".format(i)) for i in xrange(cfg.FPN.ROI_MIN_LEVEL, cfg.FPN.ROI_MAX_LEVEL + 1)] + [core.ScopedName("track_rois_idx_restore_int32"), str(blob_track)]
)
model.AddParams([core.BlobReference(input_name) for op in model.track_net.Proto().op for input_name in op.input if input_name[-2] == "_"])
del track_net_temp
else:
model.track_net, _ = c2_utils.SuffixNet(
'track_net', model.net, len(bbox_net.op), blob_track
)
model.net._net = bbox_net
loss_gradients = None
else:
loss_gradients = track_rcnn_heads.add_track_losses(model)
return loss_gradients, blob_track
