text: # --------------------------------------------------------
  # Multitask Network Cascade
  # Written by Haozhi Qi
  # Copyright (c) 2016, Haozhi Qi
  # Licensed under The MIT License [see LICENSE for details]
  # --------------------------------------------------------
  
  import caffe
  import numpy as np
  import yaml
  from transform.bbox_transform import \
      bbox_transform_inv, bbox_compute_targets, \
      clip_boxes, get_bbox_regression_label
  from transform.mask_transform import intersect_mask
  from mnc_config import cfg
  from utils.cython_bbox import bbox_overlaps
  
  
  class StageBridgeLayer(caffe.Layer):
      """
      This layer take input from bounding box prediction
      and output a set of new rois after applying transformation
      It will also provide mask/bbox regression targets
      during training phase
      """
      def setup(self, bottom, top):
          layer_params = yaml.load(self.param_str_)
          # bottom 0 is ~ n ROIs to train Fast RCNN
          # bottom 1 is ~ n * 4(1+c) bbox prediction
          # bottom 2 is ~ n * (1+c) bbox scores (seg classification)
          self._phase = str(self.phase)
          if self._phase == 'TRAIN':
              self._use_clip = layer_params['use_clip']
              self._clip_denominator = float(layer_params.get('clip_base', 64))
              self._clip_thresh = 1.0 / self._clip_denominator
              self._feat_stride = layer_params['feat_stride']
              self._num_classes = layer_params['num_classes']
  
          # meaning of top blobs speak for themselves
          self._top_name_map = {}
          if self._phase == 'TRAIN':
              top[0].reshape(1, 5)
              self._top_name_map['rois'] = 0
              top[1].reshape(1, 1)
              self._top_name_map['labels'] = 1
              top[2].reshape(1, 1, cfg.MASK_SIZE, cfg.MASK_SIZE)
              self._top_name_map['mask_targets'] = 2
              top[3].reshape(1, 1, cfg.MASK_SIZE, cfg.MASK_SIZE)
              self._top_name_map['mask_weight'] = 3
              top[4].reshape(1, 4)
              self._top_name_map['gt_mask_info'] = 4
              top[5].reshape(1, self._num_classes * 4)
              self._top_name_map['bbox_targets'] = 5
              top[6].reshape(1, self._num_classes * 4)
              self._top_name_map['bbox_inside_weights'] = 6
              top[7].reshape(1, self._num_classes * 4)
              self._top_name_map['bbox_outside_weights'] = 7
          elif self._phase == 'TEST':
              top[0].reshape(1, 5)
              self._top_name_map['rois'] = 0
          else:
              print 'Unrecognized phase'
              raise NotImplementedError
  
      def reshape(self, bottom, top):
          # reshape happens during forward
          pass
  
      def forward(self, bottom, top):
          if str(self.phase) == 'TRAIN':
              blobs = self.forward_train(bottom, top)
          elif str(self.phase) == 'TEST':
              blobs = self.forward_test(bottom, top)
          else:
              print 'Unrecognized phase'
              raise NotImplementedError
  
          for blob_name, blob in blobs.iteritems():
              top[self._top_name_map[blob_name]].reshape(*blob.shape)
              top[self._top_name_map[blob_name]].data[...] = blob.astype(np.float32, copy=False)
  
      def backward(self, top, propagate_down, bottom):
          """
          Description:
              We need to implement bp for 2 bottoms:
              The top diff is x_new, y_new, w_new, h_new
          """
          deltas = bottom[1].data
          dfdxc = top[0].diff[:, 1]
          dfdyc = top[0].diff[:, 2]
          dfdw = top[0].diff[:, 3]
          dfdh = top[0].diff[:, 4]
          W_old = bottom[0].data[:, 2] - bottom[0].data[:, 0]
          H_old = bottom[0].data[:, 3] - bottom[0].data[:, 1]
  
          if propagate_down[0]:
              bottom[0].diff.fill(0.)
              for ind, i in enumerate(self._keep_inds):
                  if i >= bottom[0].diff.shape[0] or self._bbox_reg_labels[i] == 0:
                      continue
                  delta_x = deltas[i, 4*self._bbox_reg_labels[i]]
                  delta_y = deltas[i, 4*self._bbox_reg_labels[i]+1]
                  delta_w = deltas[i, 4*self._bbox_reg_labels[i]+2]
                  delta_h = deltas[i, 4*self._bbox_reg_labels[i]+3]
                  bottom[0].diff[i, 1] = dfdxc[ind]
                  bottom[0].diff[i, 2] = dfdyc[ind]
                  bottom[0].diff[i, 3] = dfdw[ind] * (delta_x + np.exp(delta_w))
                  bottom[0].diff[i, 4] = dfdh[ind] * (delta_y + np.exp(delta_h))
  
          if propagate_down[1]:
              bottom[1].diff.fill(0.)
              for ind, i in enumerate(self._keep_inds):
                  if i >= bottom[1].diff.shape[0] or i not in self._clip_keep or self._bbox_reg_labels[i] == 0:
                      continue
                  delta_w = deltas[i, 4*self._bbox_reg_labels[i]+2]
                  delta_h = deltas[i, 4*self._bbox_reg_labels[i]+3]
                  bottom[1].diff[i, 4*self._bbox_reg_labels[i]] = dfdxc[ind] * W_old[i]
                  bottom[1].diff[i, 4*self._bbox_reg_labels[i]+1] = dfdyc[ind] * H_old[i]
                  bottom[1].diff[i, 4*self._bbox_reg_labels[i]+2] = dfdw[ind] * np.exp(delta_w) * W_old[i]
                  bottom[1].diff[i, 4*self._bbox_reg_labels[i]+3] = dfdh[ind] * np.exp(delta_h) * H_old[i]
                  if self._use_clip:
                      bottom[1].diff[i, 4*self._bbox_reg_labels[i]] = np.minimum(np.maximum(
                          bottom[1].diff[i, 4*self._bbox_reg_labels[i]], -self._clip_thresh), self._clip_thresh)
                      bottom[1].diff[i, 4*self._bbox_reg_labels[i]+1] = np.minimum(np.maximum(
                          bottom[1].diff[i, 4*self._bbox_reg_labels[i]+1], -self._clip_thresh), self._clip_thresh)
                      bottom[1].diff[i, 4*self._bbox_reg_labels[i]+2] = np.minimum(np.maximum(
                          bottom[1].diff[i, 4*self._bbox_reg_labels[i]+2], -self._clip_thresh), self._clip_thresh)
                      bottom[1].diff[i, 4*self._bbox_reg_labels[i]+3] = np.minimum(np.maximum(
                          bottom[1].diff[i, 4*self._bbox_reg_labels[i]+3], -self._clip_thresh), self._clip_thresh)
  
      def forward_train(self, bottom, top):
          """
          During forward, we need to do several things:
          1. Apply bounding box regression output which has highest
             classification score to proposed ROIs
          2. Sample ROIs based on there current overlaps, assign labels
             on them
          3. Make mask regression targets and positive/negative weights,
             just like the proposal_target_layer
          """
          rois = bottom[0].data
          bbox_deltas = bottom[1].data
          # Apply bounding box regression according to maximum segmentation score
          seg_scores = bottom[2].data
          self._bbox_reg_labels = seg_scores[:, 1:].argmax(axis=1) + 1
  
          gt_boxes = bottom[3].data
          gt_masks = bottom[4].data
          im_info = bottom[5].data[0, :]
          mask_info = bottom[6].data
  
          # select bbox_deltas according to
          artificial_deltas = np.zeros((rois.shape[0], 4))
          for i in xrange(rois.shape[0]):
              artificial_deltas[i, :] = bbox_deltas[i, 4*self._bbox_reg_labels[i]:4*(self._bbox_reg_labels[i]+1)]
          artificial_deltas[self._bbox_reg_labels == 0, :] = 0
  
          all_rois = np.zeros((rois.shape[0], 5))
          all_rois[:, 0] = 0
          all_rois[:, 1:5] = bbox_transform_inv(rois[:, 1:5], artificial_deltas)
          zeros = np.zeros((gt_boxes.shape[0], 1), dtype=gt_boxes.dtype)
          all_rois = np.vstack(
              (all_rois, np.hstack((zeros, gt_boxes[:, :-1])))
          )
          all_rois[:, 1:5], self._clip_keep = clip_boxes(all_rois[:, 1:5], im_info[:2])
  
          labels, rois_out, fg_inds, keep_inds, mask_targets, top_mask_info, bbox_targets, bbox_inside_weights = \
              self._sample_output(all_rois, gt_boxes, im_info[2], gt_masks, mask_info)
          bbox_outside_weights = np.array(bbox_inside_weights > 0).astype(np.float32)
          self._keep_inds = keep_inds
  
          mask_weight = np.zeros((rois_out.shape[0], 1, cfg.MASK_SIZE, cfg.MASK_SIZE))
          mask_weight[0:len(fg_inds), :, :, :] = 1
  
          blobs = {
              'rois': rois_out,
              'labels': labels,
              'mask_targets': mask_targets,
              'mask_weight': mask_weight,
              'gt_mask_info': top_mask_info,
              'bbox_targets': bbox_targets,
              'bbox_inside_weights': bbox_inside_weights,
              'bbox_outside_weights': bbox_outside_weights
          }
          return blobs
  
      def _sample_output(self, all_rois, gt_boxes, im_scale, gt_masks, mask_info):
          overlaps = bbox_overlaps(
              np.ascontiguousarray(all_rois[:, 1:5], dtype=np.float),
              np.ascontiguousarray(gt_boxes[:, :4], dtype=np.float))
          gt_assignment = overlaps.argmax(axis=1)
          max_overlaps = overlaps.max(axis=1)
          labels = gt_boxes[gt_assignment, 4]
          # Sample foreground indexes
          fg_inds = np.where(max_overlaps >= cfg.TRAIN.BBOX_THRESH)[0]
          bg_inds = np.where(max_overlaps < cfg.TRAIN.BBOX_THRESH)[0]
          keep_inds = np.append(fg_inds, bg_inds).astype(int)
          # Select sampled values from various arrays:
          labels = labels[keep_inds]
          # Clamp labels for the background RoIs to 0
          labels[len(fg_inds):] = 0
          rois = all_rois[keep_inds]
  
          bbox_target_data = bbox_compute_targets(
              rois[:, 1:5], gt_boxes[gt_assignment[keep_inds], :4], normalize=True)
          bbox_target_data = np.hstack((labels[:, np.newaxis], bbox_target_data))\
              .astype(np.float32, copy=False)
          bbox_targets, bbox_inside_weights = get_bbox_regression_label(
              bbox_target_data, self._num_classes)
  
          scaled_rois = rois[:, 1:5] / float(im_scale)
          scaled_gt_boxes = gt_boxes[:, :4] / float(im_scale)
  
          pos_masks = np.zeros((len(keep_inds), 1,  cfg.MASK_SIZE,  cfg.MASK_SIZE))
          top_mask_info = np.zeros((len(keep_inds), 12))
          top_mask_info[len(fg_inds):, :] = -1
  
          for i, val in enumerate(fg_inds):
              gt_box = scaled_gt_boxes[gt_assignment[val]]
              gt_box = np.around(gt_box).astype(int)
              ex_box = np.around(scaled_rois[i]).astype(int)
              gt_mask = gt_masks[gt_assignment[val]]
              gt_mask_info = mask_info[gt_assignment[val]]
              gt_mask = gt_mask[0:gt_mask_info[0], 0:gt_mask_info[1]]
              # regression targets is the intersection of bounding box and gt mask
              ex_mask = intersect_mask(ex_box, gt_box, gt_mask)
              pos_masks[i, ...] = ex_mask
              top_mask_info[i, 0] = gt_assignment[val]
              top_mask_info[i, 1] = gt_mask_info[0]
              top_mask_info[i, 2] = gt_mask_info[1]
              top_mask_info[i, 3] = labels[i]
              top_mask_info[i, 4:8] = ex_box
              top_mask_info[i, 8:12] = gt_box
  
          return labels, rois, fg_inds, keep_inds, pos_masks, top_mask_info, bbox_targets, bbox_inside_weights
  
      def forward_test(self, bottom, top):
          rois = bottom[0].data
          bbox_deltas = bottom[1].data
          # get ~ n * 4(1+c) new rois
          all_rois = bbox_transform_inv(rois[:, 1:5], bbox_deltas)
          scores = bottom[2].data
          im_info = bottom[3].data
          # get highest scored category's bounding box regressor
          score_max = scores.argmax(axis=1)
          rois_out = np.zeros((rois.shape[0], 5))
          # Single batch training
          rois_out[:, 0] = 0
          for i in xrange(len(score_max)):
              rois_out[i, 1:5] = all_rois[i, 4*score_max[i]:4*(score_max[i]+1)]
          rois_out[:, 1:5], _ = clip_boxes(rois_out[:, 1:5], im_info[0, :2])
          blobs = {
              'rois': rois_out
          }
          return blobs
  