def add_body_uv_losses()

in detectron/modeling/body_uv_rcnn_heads.py [0:0]


def add_body_uv_losses(model, pref=''):

    ## Reshape for GT blobs.
    model.net.Reshape( ['body_uv_X_points'], ['X_points_reshaped'+pref, 'X_points_shape'+pref],  shape=( -1 ,1 ) )
    model.net.Reshape( ['body_uv_Y_points'], ['Y_points_reshaped'+pref, 'Y_points_shape'+pref],  shape=( -1 ,1 ) )
    model.net.Reshape( ['body_uv_I_points'], ['I_points_reshaped'+pref, 'I_points_shape'+pref],  shape=( -1 ,1 ) )
    model.net.Reshape( ['body_uv_Ind_points'], ['Ind_points_reshaped'+pref, 'Ind_points_shape'+pref],  shape=( -1 ,1 ) )
    ## Concat Ind,x,y to get Coordinates blob.
    model.net.Concat( ['Ind_points_reshaped'+pref,'X_points_reshaped'+pref, \
                       'Y_points_reshaped'+pref],['Coordinates'+pref,'Coordinate_Shapes'+pref ], axis = 1 )
    ##
    ### Now reshape UV blobs, such that they are 1x1x(196*NumSamples)xNUM_PATCHES 
    ## U blob to
    ##
    model.net.Reshape(['body_uv_U_points'], \
                      ['U_points_reshaped'+pref, 'U_points_old_shape'+pref],\
                      shape=(-1,cfg.BODY_UV_RCNN.NUM_PATCHES+1,196))
    model.net.Transpose(['U_points_reshaped'+pref] ,['U_points_reshaped_transpose'+pref],axes=(0,2,1) )
    model.net.Reshape(['U_points_reshaped_transpose'+pref], \
                      ['U_points'+pref, 'U_points_old_shape2'+pref], \
                      shape=(1,1,-1,cfg.BODY_UV_RCNN.NUM_PATCHES+1))
    ## V blob
    ##
    model.net.Reshape(['body_uv_V_points'], \
                      ['V_points_reshaped'+pref, 'V_points_old_shape'+pref],\
                      shape=(-1,cfg.BODY_UV_RCNN.NUM_PATCHES+1,196))
    model.net.Transpose(['V_points_reshaped'+pref] ,['V_points_reshaped_transpose'+pref],axes=(0,2,1) )
    model.net.Reshape(['V_points_reshaped_transpose'+pref], \
                      ['V_points'+pref, 'V_points_old_shape2'+pref], \
                      shape=(1,1,-1,cfg.BODY_UV_RCNN.NUM_PATCHES+1))
    ###
    ## UV weights blob
    ##
    model.net.Reshape(['body_uv_point_weights'], \
                      ['Uv_point_weights_reshaped'+pref, 'Uv_point_weights_old_shape'+pref],\
                      shape=(-1,cfg.BODY_UV_RCNN.NUM_PATCHES+1,196))
    model.net.Transpose(['Uv_point_weights_reshaped'+pref] ,['Uv_point_weights_reshaped_transpose'+pref],axes=(0,2,1) )
    model.net.Reshape(['Uv_point_weights_reshaped_transpose'+pref], \
                      ['Uv_point_weights'+pref, 'Uv_point_weights_old_shape2'+pref], \
                      shape=(1,1,-1,cfg.BODY_UV_RCNN.NUM_PATCHES+1))

    #####################
    ###  Pool IUV for points via bilinear interpolation.
    model.PoolPointsInterp(['U_estimated','Coordinates'+pref], ['interp_U'+pref])
    model.PoolPointsInterp(['V_estimated','Coordinates'+pref], ['interp_V'+pref])
    model.PoolPointsInterp(['Index_UV'+pref,'Coordinates'+pref], ['interp_Index_UV'+pref])

    ## Reshape interpolated UV coordinates to apply the loss.
    
    model.net.Reshape(['interp_U'+pref], \
                      ['interp_U_reshaped'+pref, 'interp_U_shape'+pref],\
                      shape=(1, 1, -1 , cfg.BODY_UV_RCNN.NUM_PATCHES+1))
    
    model.net.Reshape(['interp_V'+pref], \
                      ['interp_V_reshaped'+pref, 'interp_V_shape'+pref],\
                      shape=(1, 1, -1 , cfg.BODY_UV_RCNN.NUM_PATCHES+1))
    ###

    ### Do the actual labels here !!!!
    model.net.Reshape( ['body_uv_ann_labels'],    \
                      ['body_uv_ann_labels_reshaped'   +pref, 'body_uv_ann_labels_old_shape'+pref], \
                      shape=(-1, cfg.BODY_UV_RCNN.HEATMAP_SIZE , cfg.BODY_UV_RCNN.HEATMAP_SIZE))
    
    model.net.Reshape( ['body_uv_ann_weights'],   \
                      ['body_uv_ann_weights_reshaped'   +pref, 'body_uv_ann_weights_old_shape'+pref], \
                      shape=( -1 , cfg.BODY_UV_RCNN.HEATMAP_SIZE , cfg.BODY_UV_RCNN.HEATMAP_SIZE))
    ###
    model.net.Cast( ['I_points_reshaped'+pref], ['I_points_reshaped_int'+pref], to=core.DataType.INT32)
    ### Now add the actual losses 
    ## The mask segmentation loss (dense)
    probs_seg_AnnIndex, loss_seg_AnnIndex = model.net.SpatialSoftmaxWithLoss( \
                          ['AnnIndex'+pref, 'body_uv_ann_labels_reshaped'+pref,'body_uv_ann_weights_reshaped'+pref],\
                          ['probs_seg_AnnIndex'+pref,'loss_seg_AnnIndex'+pref], \
                           scale=cfg.BODY_UV_RCNN.INDEX_WEIGHTS / cfg.NUM_GPUS)
    ## Point Patch Index Loss.
    probs_IndexUVPoints, loss_IndexUVPoints = model.net.SoftmaxWithLoss(\
                          ['interp_Index_UV'+pref,'I_points_reshaped_int'+pref],\
                          ['probs_IndexUVPoints'+pref,'loss_IndexUVPoints'+pref], \
                          scale=cfg.BODY_UV_RCNN.PART_WEIGHTS / cfg.NUM_GPUS, spatial=0)
    ## U and V point losses.
    loss_Upoints = model.net.SmoothL1Loss( \
                          ['interp_U_reshaped'+pref, 'U_points'+pref, \
                               'Uv_point_weights'+pref, 'Uv_point_weights'+pref], \
                          'loss_Upoints'+pref, \
                            scale=cfg.BODY_UV_RCNN.POINT_REGRESSION_WEIGHTS  / cfg.NUM_GPUS)
    
    loss_Vpoints = model.net.SmoothL1Loss( \
                          ['interp_V_reshaped'+pref, 'V_points'+pref, \
                               'Uv_point_weights'+pref, 'Uv_point_weights'+pref], \
                          'loss_Vpoints'+pref, scale=cfg.BODY_UV_RCNN.POINT_REGRESSION_WEIGHTS / cfg.NUM_GPUS)
    ## Add the losses.
    loss_gradients = blob_utils.get_loss_gradients(model, \
                       [ loss_Upoints, loss_Vpoints, loss_seg_AnnIndex, loss_IndexUVPoints])
    model.losses = list(set(model.losses + \
                       ['loss_Upoints'+pref , 'loss_Vpoints'+pref , \
                        'loss_seg_AnnIndex'+pref ,'loss_IndexUVPoints'+pref]))

    return loss_gradients