def _add_roi_keypoint_head( model, add_roi_keypoint_head_func, blob_in, dim_in, spatial_scale_in ): """Add a keypoint prediction head to the model.""" # Capture model graph before adding the mask head bbox_net = copy.deepcopy(model.net.Proto()) # Add the keypoint head blob_keypoint_head, dim_keypoint_head = add_roi_keypoint_head_func( model, blob_in, dim_in, spatial_scale_in ) # Add the keypoint output blob_keypoint = keypoint_rcnn_heads.add_keypoint_outputs( model, blob_keypoint_head, dim_keypoint_head ) if not model.train: # == inference # Inference uses a cascade of box predictions, then keypoint predictions # This requires separate nets for box and keypoint 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.keypoint_net, keypoint_blob_out = c2_utils.SuffixNet( 'keypoint_net', model.net, len(bbox_net.op), blob_keypoint ) model.net._net = bbox_net loss_gradients = None else: loss_gradients = keypoint_rcnn_heads.add_keypoint_losses(model) return loss_gradients
def _add_roi_boundary_head(model, add_roi_boundary_head_func, blob_in, dim_in, spatial_scale_in): """Add a boundary 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_boundary_head, dim_boundary_head = add_roi_boundary_head_func( model, blob_in, dim_in, spatial_scale_in) # Add the mask output blob_boundary = boundary_heads.add_boundary_rcnn_outputs( model, blob_boundary_head, dim_boundary_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.boundary_net, blob_boundary = c2_utils.SuffixNet( 'boundary_net', model.net, len(bbox_net.op), blob_boundary) model.net._net = bbox_net loss_gradients = None else: loss_gradients = boundary_heads.add_boundary_rcnn_losses( model, blob_boundary) return loss_gradients, blob_boundary_head
def _add_semantic_segms_head(model, blob_in, dim_in, spatial_scale_in): """ Add the semantic segmentation head to the network """ conv_body_net = copy.deepcopy(model.net.Proto()) # Add semantic segms net's input rescale_factor = cfg.SEMANTIC_NET.RESCALE_FACTOR blob_rescale_feat, dim_rescale_feat = model.RescaleFeatureMap( blob_in, 'img_rescale_feature', dim_in, rescale_factor=cfg.SEMANTIC_NET.RESCALE_FACTOR, spatial_scale=spatial_scale_in, sampling_ratio=cfg.SEMANTIC_NET.ROI_XFORM_SAMPLING_RATIO) # Add semantic segms net head blob_semantic_segms_head, dim_semantic_segms_head = \ semantic_segms_head.add_semantic_segms_head( model, blob_rescale_feat, dim_rescale_feat, ) # Add semantic segms net output blob_semantic_segms_out = semantic_segms_head.add_semantic_segms_outputs( model, blob_semantic_segms_head, dim_semantic_segms_head) # Add semantic segms loss if not model.train: # inference model.semantic_segms_net, semantic_segms_out = c2_utils.SuffixNet( 'semantic_segms_net', model.net, len(conv_body_net.op), blob_semantic_segms_out) model.net._net = conv_body_net loss_gradients = None else: loss_gradients = semantic_segms_head.add_semantic_segms_losses( model, blob_semantic_segms_out) return loss_gradients
def _add_prn_head(model, blob_in, dim_in, spatial_scale_in): """ Add a classification head to predict whether the roi needs further refinement prn mean predict refinement-needed """ # Capture the model graph before adding the prn head bbox_net = copy.deepcopy(model.net.Proto()) # First generate labels for prn head and update blobs for RefineNet if model.train: prn_heads.add_prn_labels(model) # Add the prediction heads prefix = 'mask' blob_prn_head, dim_prn_head = prn_heads.add_prn_head( model, blob_in, dim_in, spatial_scale_in, prefix) # Add the prediction output blob_prn_out = prn_heads.add_prn_outputs(model, blob_prn_head, dim_prn_head) if not model.train: # == inference # Inference uses a cascade of box predictions, then roi mask predictions # then refine mask prediction. This requires separate nets for box and # mask and refine mask prediction. # So we extract the need refinement prediction net, store it as its own # network,then restore model.net to be the bbox-only network model.prn_net, prn_blob_out = c2_utils.SuffixNet( 'prn_net', model.net, len(bbox_net.op), blob_prn_out) model.net._net = bbox_net loss_gradients = None else: loss_gradients = prn_heads.add_prn_losses(model) return loss_gradients
def _add_generic_refine_keypoint_net_head(model, blob_in, dim_in, spatial_scale_in): """Add a generic refine head to the model. Allows Mask/Keypoint indicator and Mask/Keypoint refined output """ # Capture the model graph before adding the refine head bbox_net = copy.deepcopy(model.net.Proto()) # Prepare for RefineNet input # Different indicator type will generate different blob_refine_net_in INDICATOR_TYPE = cfg.REFINENET.INDICATOR_TYPE blob_refine_net_in, dim_refine_net_in = \ refine_net_heads.add_refine_net_inputs( model, blob_in, dim_in, spatial_scale_in, INDICATOR_TYPE ) # Add RefineNet head prefix = 'keypoint' blob_refine_net_head, dim_refine_net_head = \ refine_net_heads.add_refine_net_head( model, blob_refine_net_in, dim_refine_net_in, prefix ) # Add RefineNet output # Different refined-output type will generate different output blob_refine_keypoint_out = refine_net_heads.add_refine_net_keypoint_outputs( model, blob_refine_net_head, dim_refine_net_head) if not model.train: # == inference # Inference uses a cascade of box predictions, then roi mask predictions # then refine mask prediction. This requires separate nets for box and # mask and refine mask prediction. # So we extract the refine mask prediction net, store it as its own # network,then restore model.net to be the bbox-only network model.refine_keypoint_net, refine_blob_out = c2_utils.SuffixNet( 'refine_keypoint_net', model.net, len(bbox_net.op), blob_refine_keypoint_out) model.net._net = bbox_net loss_gradients = None else: loss_gradients = refine_net_heads.add_refine_net_keypoint_losses( model, blob_refine_keypoint_out) return loss_gradients
def _add_roi_mask_head(model, add_roi_mask_head_func, blob_in, dim_in, spatial_scale_in, blob_boundary_attention): """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) # attention if cfg.MODEL.BOUNDARY_ON and cfg.BOUNDARY.CONCAT_MASK: # s = model.net.Sum([blob_mask_head, blob_boundary_attention], 'm_b_sum') bo = 'm_b_concat' concat_F, concat_dims = model.net.Concat( [blob_mask_head, blob_boundary_attention], [bo, "_" + bo + "_concat_dims"]) concat_dims = dim_mask_head * 2 else: concat_F = blob_mask_head concat_dims = dim_mask_head # Add the mask output blob_mask = mask_rcnn_heads.add_mask_rcnn_outputs(model, concat_F, concat_dims) 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