SEGMENTATION_LABELS = "segmentation_labels"
SEGMENTATION_LABELS_ORIGINAL_SIZE = "segmentation_labels_original_size"
BBOX_GUIDANCE = "bbox_guidance"
UNSIGNED_DISTANCE_TRANSFORM_GUIDANCE = "unsigned_distance_transform_guidance"
SIGNED_DISTANCE_TRANSFORM_GUIDANCE = "signed_distance_transform_guidance"
LASER_GUIDANCE = "laser_guidance"
IMAGES = "images"
RAW_IMAGES = "raw_images"
INPUTS = "inputs"
IMAGE_FILENAMES = "image_filenames"
CLASSES = "classes"
IDS = "ids"  # bounding box number 1,2,3...  or 0 if no bounding box
DT_METHOD = "dt_method"
RAW_SEGMENTATION_LABELS = "raw_segmentation_labels"
CLICK_MAPS = "clicks_maps"
NEG_CLICKS = "neg_clicks"
POS_CLICKS = "pos_clicks"
OBJ_TAGS = "obj_tags"
FEATUREMAP_LABELS = "featuremap_labels"
FEATUREMAP_BOXES = "featuremap_boxes"
IMAGE_ID = "image_id"
IS_CROWD = "is_crowd"
SEGMENTATION_MASK = "segmentation_mask"
SKIP_EXAMPLE = "skip_example"  # If true, filter the corresponding example from the datastream. Useful, when an example
RAW_IMAGE_SIZES = "raw_image_sizes"  # h, w
TIME_OFFSETS = "time_offsets"
SPACE_OFFSETS = "space_offsets"
SEGMENTATION_INSTANCE_LABELS = "segmentation_instance_labels"
N_VALID_IDS = "n_valid_ids" class Conv(Layer):
  output_layer = False

  def __init__(self, name, inputs, n_features, tower_setup, filter_size=(3, 3), old_order=False,
               strides=(1, 1), dilation=None, pool_size=(1, 1), pool_strides=None, activation="relu", dropout=0.0,
               batch_norm=False, bias=False, batch_norm_decay=Layer.BATCH_NORM_DECAY_DEFAULT, l2=Layer.L2_DEFAULT,
               padding="SAME", W_initializer=None, b_initializer=None, freeze_batchnorm_override=None):
    super(Conv, self).__init__()
    # mind the order of dropout, conv, activation and batchnorm!
    # batchnorm -> activation -> dropout -> conv -> pool
    # if old_order: dropout -> conv -> batchnorm -> activation -> pool (used for example in tensorpack)
    curr, n_features_inp = prepare_input(inputs)
    filter_size = list(filter_size)
    strides = list(strides)
    pool_size = list(pool_size)
    if pool_strides is None:
      pool_strides = pool_size

    with tf.variable_scope(name):
      if isinstance(W_initializer, list):
        W_initializer = tf.constant_initializer(W_initializer)
      W = self.create_weight_variable("W", filter_size + [n_features_inp, n_features], l2, tower_setup,
                                      initializer=W_initializer)
      b = None
      if bias:
        if b_initializer is not None:
          b_initializer = tf.constant_initializer(b_initializer)
          b = self.create_bias_variable("b", [n_features], tower_setup, initializer=b_initializer)
        else:
          b = self.create_bias_variable("b", [n_features], tower_setup)

      if old_order:
        curr = apply_dropout(curr, dropout)
        if dilation is None:
          curr = conv2d(curr, W, strides, padding=padding)
        else:
          curr = conv2d_dilated(curr, W, dilation, padding=padding)
        if bias:
          curr += b
        if batch_norm:
          curr = self.create_and_apply_batch_norm(curr, n_features, batch_norm_decay, tower_setup, freeze_batchnorm_override=freeze_batchnorm_override)
        curr = get_activation(activation)(curr)
      else:
        if batch_norm:
          curr = self.create_and_apply_batch_norm(curr, n_features_inp, batch_norm_decay, tower_setup, freeze_batchnorm_override=freeze_batchnorm_override)
        curr = get_activation(activation)(curr)
        curr = apply_dropout(curr, dropout)
        if dilation is None:
          curr = conv2d(curr, W, strides, padding=padding)
        else:
          curr = conv2d_dilated(curr, W, dilation, padding=padding)
        if bias:
          curr += b

      if pool_size != [1, 1]:
        curr = max_pool(curr, pool_size, pool_strides)
    self.outputs = [curr] class ConvTranspose(Layer):
  output_layer = False

  def __init__(self, name, inputs, n_features, tower_setup, filter_size=(3, 3),
               strides=(1, 1), activation="relu",  # dropout=0.0,
               batch_norm=False, bias=False, batch_norm_decay=Layer.BATCH_NORM_DECAY_DEFAULT, l2=Layer.L2_DEFAULT,
               padding="SAME", W_initializer=None, freeze_batchnorm_override=None):
    super(ConvTranspose, self).__init__()
    # TODO this uses the tensorpack order of operations
    curr, n_features_inp = prepare_input(inputs)
    filter_size = list(filter_size)
    strides = list(strides)

    with tf.variable_scope(name):
      W = self.create_weight_variable("W", filter_size + [n_features, n_features_inp], l2, tower_setup,
                                      initializer=W_initializer)
      b = None
      if bias:
        b = self.create_bias_variable("b", [n_features], tower_setup)

      # curr = apply_dropout(curr, dropout)
      curr = conv2d_transpose(curr, W, strides, padding=padding)
      if bias:
        curr = tf.nn.bias_add(curr, b)
      if batch_norm:
        curr = self.create_and_apply_batch_norm(curr, n_features, batch_norm_decay, tower_setup,
                                                freeze_batchnorm_override=freeze_batchnorm_override)
      curr = get_activation(activation)(curr)
    self.outputs = [curr] class FullyConnected(Layer):
  def __init__(self, name, inputs, n_features, tower_setup, activation="relu", dropout=0.0, batch_norm=False,
               batch_norm_decay=Layer.BATCH_NORM_DECAY_DEFAULT, l2=Layer.L2_DEFAULT, W_initializer=None):
    super(FullyConnected, self).__init__()
    inp, n_features_inp = prepare_input(inputs)

    with tf.variable_scope(name):
      inp = apply_dropout(inp, dropout)
      if batch_norm:
        inp = tf.expand_dims(inp, axis=0)
        inp = tf.expand_dims(inp, axis=0)
        inp = self.create_and_apply_batch_norm(inp, n_features_inp, batch_norm_decay, tower_setup)
        inp = tf.squeeze(inp, axis=[0, 1])
      W = self.create_weight_variable("W", [n_features_inp, n_features], l2, tower_setup, initializer=W_initializer)
      b = self.create_bias_variable("b", [n_features], tower_setup)
      z = tf.matmul(inp, W) + b
      h = get_activation(activation)(z)
    self.outputs = [h]
    self.n_features = n_features class Layer:
  BATCH_NORM_DECAY_DEFAULT = 0.95
  BATCH_NORM_EPSILON = 1e-5
  L2_DEFAULT = 1e-4

  def __init__(self):
    self.summaries = []
    self.regularizers = []
    self.losses = []
    self.update_ops = []
    self.outputs = []
    self.placeholders = []
    self.measures = {}
    self.extractions = {}
    self.n_params = 0

  def add_scalar_summary(self, op, name):
    summary = tf.summary.scalar(name, op)
    self.summaries.append(summary)

  def add_image_summary(self, im, name):
    summary = tf.summary.image(name, im)
    self.summaries.append(summary)

  def create_and_apply_batch_norm(self, inp, n_features, decay, tower_setup, scope_name="bn", freeze_batchnorm_override=None):
    beta, gamma, moving_mean, moving_var = create_batch_norm_vars(n_features, tower_setup, scope_name)
    self.n_params += 2 * n_features
    if tower_setup.is_main_train_tower:
      assert tower_setup.is_training
    if tower_setup.is_training and\
      (not tower_setup.freeze_batchnorm or (freeze_batchnorm_override is not None and freeze_batchnorm_override)):
        xn, batch_mean, batch_var = tf.nn.fused_batch_norm(inp, gamma, beta, epsilon=Layer.BATCH_NORM_EPSILON,
                                                           is_training=True)
        if tower_setup.is_main_train_tower:
          update_op1 = moving_averages.assign_moving_average(
            moving_mean, batch_mean, decay, zero_debias=False, name='mean_ema_op')
          update_op2 = moving_averages.assign_moving_average(
            moving_var, batch_var, decay, zero_debias=False, name='var_ema_op')
          self.update_ops.append(update_op1)
          self.update_ops.append(update_op2)
        return xn
    else:
      # According to the tensorpack code, this updates gamma and beta but not the moving averages
      # Fused version should have better performance.
      xn, _, _ = tf.nn.fused_batch_norm(inp, gamma, beta, moving_mean, moving_var, Layer.BATCH_NORM_EPSILON, is_training=False)
      return xn

  def create_weight_variable(self, name, shape, l2, tower_setup, trainable=True, initializer=None):
    with tf.device(tower_setup.variable_device):
      if initializer is None:
        # He initialization
        initializer = tf.contrib.layers.variance_scaling_initializer(factor=2.0, mode='FAN_IN', uniform=False)
      self.n_params += np.prod(shape)
      if type(initializer) == np.ndarray:
        W = tf.get_variable(name, dtype=tower_setup.dtype, initializer=initializer, trainable=trainable)
      else:
        W = tf.get_variable(name, shape, tower_setup.dtype, initializer, trainable=trainable)
      if l2 > 0.0:
        self.regularizers.append(l2 * tf.nn.l2_loss(W))
      if tower_setup.use_weight_summaries:
        summ = tf.summary.histogram(name, W)
        self.summaries.append(summ)
        self.add_scalar_summary(tf.reduce_max(tf.abs(W)), name + "/W_abs_max")
      return W

  def create_bias_variable(self, name, shape, tower_setup, trainable=True, initializer=None):
    with tf.device(tower_setup.variable_device):
      if initializer is None:
        initializer = tf.constant_initializer(0.0, dtype=tower_setup.dtype)
      self.n_params += np.prod(shape)
      b = tf.get_variable(name, shape, tower_setup.dtype, initializer, trainable=trainable)
      if tower_setup.use_weight_summaries:
        summ = tf.summary.histogram(name, b)
        self.summaries.append(summ)
        self.add_scalar_summary(tf.reduce_max(tf.abs(b)), name + "/b_abs_max")
      return b def add_resnet_conv5(image, tower_setup):
  c5 = resnet_group(image, 'group3', 512, 3, 2, tower_setup)
  return [c5] VOID_LABEL = 255
MEASURES = "measures"
N_EXAMPLES = "n_examples"
LOSS = "loss"
IOU = "IoU"
RECALL = "recall"
PRECISION = "precision"
MAP_BBOX = "mAP (bbox)"
MAP_BBOX_50 = "mAP (bbox) @50%"
MAP_SEGM = "mAP (segm)"
ERROR_RATE = "error_rate"
ERROR_RATE_AT_005 = "error_rate_at_005"
CONFUSION_MATRIX = "confusion_matrix"
  T = np.logical_and(target, valid_mask).sum()
  P = np.logical_and(prediction, valid_mask).sum()
  I = np.logical_and(prediction == 1, target == 1).sum()
  U = np.logical_and(np.logical_or(prediction == 1, target == 1), valid_mask).sum()
def accumulate_measures(measures_accumulator, *new_measures, exclude=()):
def compute_measures_average(measures, for_final_result, exclude=()):
def measures_string_to_print(measures, exclude=()):
def compute_measures_for_binary_segmentation_tf(predictions, targets):
  def f(ps, ts):
def compute_measures_for_binary_segmentation_summed(predictions, targets):
def compute_measures_for_binary_segmentation_single_image(prediction, target): EXTRACTIONS = "extractions"
CLASS_POSTERIORS = "class_posteriors"
SEGMENTATION_POSTERIORS = "segmentation_posteriors"
SEGMENTATION_MASK_ORIGINAL_SIZE = "segmentation_mask_original_size"
SEGMENTATION_MASK_INPUT_SIZE = "segmentation_mask_input_size"
RECURRENT_STATE = "recurrent_state"
REID_FEATURES = "reid_features"
DET_BOXES = "det_boxes"
DET_PROBS = "det_probs"
DET_LABELS = "det_labels"
DET_MASKS = "det_masks"
IMAGE_ID = "img_id"
CUSTOM_STRING = "custom_string"
PRED_SPACE_TIME_OFFSETS = "pred_space_time_offsets"
PIXEL_POSITIONS = "pixel_positions"
FEATUREMAP_SIZE = "featuremap_size"
def accumulate_extractions(extractions_accumulator, *new_extractions):
def extract_batch_size_1(extractions, key): def decode_bbox_target(BBOX_DECODE_CLIP, box_predictions, anchors):
  """
  Args:
      box_predictions: (..., 4), logits
      anchors: (..., 4), floatbox. Must have the same shape
  Returns:
      box_decoded: (..., 4), float32. With the same shape.
  """
  orig_shape = tf.shape(anchors)
  box_pred_txtytwth = tf.reshape(box_predictions, (-1, 2, 2))
  box_pred_txty, box_pred_twth = tf.split(box_pred_txtytwth, 2, axis=1)
  # each is (...)x1x2
  anchors_x1y1x2y2 = tf.reshape(anchors, (-1, 2, 2))
  anchors_x1y1, anchors_x2y2 = tf.split(anchors_x1y1x2y2, 2, axis=1)

  waha = anchors_x2y2 - anchors_x1y1
  xaya = (anchors_x2y2 + anchors_x1y1) * 0.5

  wbhb = tf.exp(tf.minimum(
    box_pred_twth, BBOX_DECODE_CLIP)) * waha
  xbyb = box_pred_txty * waha + xaya
  x1y1 = xbyb - wbhb * 0.5
  x2y2 = xbyb + wbhb * 0.5  # (...)x1x2
  out = tf.concat([x1y1, x2y2], axis=-2)
  return tf.reshape(out, orig_shape) def encode_bbox_target(boxes, anchors):
  """
  Args:
      boxes: (..., 4), float32
      anchors: (..., 4), float32
  Returns:
      box_encoded: (..., 4), float32 with the same shape.
  """
  anchors_x1y1x2y2 = tf.reshape(anchors, (-1, 2, 2))
  anchors_x1y1, anchors_x2y2 = tf.split(anchors_x1y1x2y2, 2, axis=1)
  waha = anchors_x2y2 - anchors_x1y1
  xaya = (anchors_x2y2 + anchors_x1y1) * 0.5

  boxes_x1y1x2y2 = tf.reshape(boxes, (-1, 2, 2))
  boxes_x1y1, boxes_x2y2 = tf.split(boxes_x1y1x2y2, 2, axis=1)
  wbhb = boxes_x2y2 - boxes_x1y1
  xbyb = (boxes_x2y2 + boxes_x1y1) * 0.5

  # Note that here not all boxes are valid. Some may be zero
  txty = (xbyb - xaya) / waha
  twth = tf.log(wbhb / waha)  # may contain -inf for invalid boxes
  encoded = tf.concat([txty, twth], axis=1)  # (-1x2x2)
  return tf.reshape(encoded, tf.shape(boxes)) def generate_rpn_proposals(boxes, scores, img_shape, is_training):
  """
  Args:
      boxes: nx4 float dtype, decoded to floatbox already
      scores: n float, the logits
      img_shape: [h, w]
  Returns:
      boxes: kx4 float
      scores: k logits
  """
  assert boxes.shape.ndims == 2, boxes.shape
  if is_training:
    PRE_NMS_TOPK = TRAIN_PRE_NMS_TOPK
    POST_NMS_TOPK = TRAIN_POST_NMS_TOPK
  else:
    PRE_NMS_TOPK = TEST_PRE_NMS_TOPK
    POST_NMS_TOPK = TEST_POST_NMS_TOPK

  topk = tf.minimum(PRE_NMS_TOPK, tf.size(scores))
  topk_scores, topk_indices = tf.nn.top_k(scores, k=topk, sorted=False)
  topk_boxes = tf.gather(boxes, topk_indices)
  topk_boxes = clip_boxes(topk_boxes, img_shape)

  topk_boxes_x1y1x2y2 = tf.reshape(topk_boxes, (-1, 2, 2))
  topk_boxes_x1y1, topk_boxes_x2y2 = tf.split(topk_boxes_x1y1x2y2, 2, axis=1)
  # nx1x2 each
  wbhb = tf.squeeze(topk_boxes_x2y2 - topk_boxes_x1y1, axis=1)
  valid = tf.reduce_all(wbhb > RPN_MIN_SIZE, axis=1)  # n,
  topk_valid_boxes_x1y1x2y2 = tf.boolean_mask(topk_boxes_x1y1x2y2, valid)
  topk_valid_scores = tf.boolean_mask(topk_scores, valid)

  topk_valid_boxes_y1x1y2x2 = tf.reshape(
    tf.reverse(topk_valid_boxes_x1y1x2y2, axis=[2]),
    (-1, 4), name='nms_input_boxes')
  # Version from tensorflow 1.10, where everything works
  from tensorflow.python.ops import gen_image_ops
  tf.image.non_max_suppression = gen_image_ops.non_max_suppression
  nms_indices = tf.image.non_max_suppression(
    topk_valid_boxes_y1x1y2x2,
    topk_valid_scores,
    max_output_size=POST_NMS_TOPK,
    iou_threshold=RPN_PROPOSAL_NMS_THRESH)

  topk_valid_boxes = tf.reshape(topk_valid_boxes_x1y1x2y2, (-1, 4))
  final_boxes = tf.gather(
    topk_valid_boxes,
    nms_indices, name='boxes')
  final_scores = tf.gather(topk_valid_scores, nms_indices, name='scores')
  tf.sigmoid(final_scores, name='probs')  # for visualization
  return final_boxes, final_scores def sample_fast_rcnn_targets(boxes, gt_boxes, gt_labels, fastrcnn_batch_per_img, target_ids):
  """
  Sample some ROIs from all proposals for training.
  Args:
      boxes: nx4 region proposals, floatbox
      gt_boxes: mx4, floatbox
      gt_labels: m, int32
  Returns:
      sampled_boxes: tx4 floatbox, the rois
      sampled_labels: t labels, in [0, #class-1]. Positive means foreground.
      fg_inds_wrt_gt: #fg indices, each in range [0, m-1].
          It contains the matching GT of each foreground roi.
  """
  iou = pairwise_iou(boxes, gt_boxes)  # nxm
  # proposal_metrics(iou)

  # add ground truth as proposals as well
  boxes = tf.concat([boxes, gt_boxes], axis=0)  # (n+m) x 4
  iou = tf.concat([iou, tf.eye(tf.shape(gt_boxes)[0])], axis=0)  # (n+m) x m

  # #proposal=n+m from now on

  def sample_fg_bg(iou):
    fg_mask = tf.reduce_max(iou, axis=1) >= FASTRCNN_FG_THRESH

    fg_inds = tf.reshape(tf.where(fg_mask), [-1])
    num_fg = tf.minimum(int(
      fastrcnn_batch_per_img * FASTRCNN_FG_RATIO),
      tf.size(fg_inds), name='num_fg')
    fg_inds = tf.random_shuffle(fg_inds)[:num_fg]

    bg_inds = tf.reshape(tf.where(tf.logical_not(fg_mask)), [-1])
    num_bg = tf.minimum(
      fastrcnn_batch_per_img - num_fg,
      tf.size(bg_inds), name='num_bg')
    bg_inds = tf.random_shuffle(bg_inds)[:num_bg]

    # add_moving_summary(num_fg, num_bg)
    return fg_inds, bg_inds

  fg_inds, bg_inds = sample_fg_bg(iou)
  # fg,bg indices w.r.t proposals

  best_iou_ind = tf.argmax(iou, axis=1)  # #proposal, each in 0~m-1
  fg_inds_wrt_gt = tf.gather(best_iou_ind, fg_inds)  # num_fg

  all_indices = tf.concat([fg_inds, bg_inds], axis=0)  # indices w.r.t all n+m proposal boxes
  ret_boxes = tf.gather(boxes, all_indices, name='sampled_proposal_boxes')

  ret_labels = tf.concat(
    [tf.gather(gt_labels, fg_inds_wrt_gt),
     tf.zeros_like(bg_inds, dtype=tf.int64)], axis=0, name='sampled_labels')

  ret_target_ids = tf.concat([tf.gather(target_ids, fg_inds_wrt_gt), tf.zeros_like(bg_inds, dtype=target_ids.dtype)],
                             axis=0, name="sampled_target_ids")
  # stop the gradient -- they are meant to be ground-truth
  return tf.stop_gradient(ret_boxes), tf.stop_gradient(ret_labels), fg_inds_wrt_gt, tf.stop_gradient(ret_target_ids) def roi_align(featuremap, boxes, output_shape):
  """
  Args:
      featuremap: 1xHxWxC
      boxes: Nx4 floatbox
      output_shape: int
  Returns:
      NxoHxoWxC
  """
  boxes = tf.stop_gradient(boxes)  # TODO
  # sample 4 locations per roi bin
  ret = crop_and_resize(
    featuremap, boxes,
    tf.zeros([tf.shape(boxes)[0]], dtype=tf.int32),
    output_shape * 2)
  ret = tf.nn.avg_pool(ret, [1, 2, 2, 1], [1, 2, 2, 1], padding='SAME', data_format='NHWC')
  return ret def rpn_losses(anchor_labels, anchor_boxes, label_logits, box_logits):
  """
  Args:
      anchor_labels: fHxfWxNA
      anchor_boxes: fHxfWxNAx4, encoded
      label_logits:  fHxfWxNA
      box_logits: fHxfWxNAx4
  Returns:
      label_loss, box_loss
  """
  with tf.device('/cpu:0'):
    valid_mask = tf.stop_gradient(tf.not_equal(anchor_labels, -1))
    pos_mask = tf.stop_gradient(tf.equal(anchor_labels, 1))
    nr_valid = tf.stop_gradient(tf.count_nonzero(valid_mask, dtype=tf.int32), name='num_valid_anchor')
    nr_pos = tf.count_nonzero(pos_mask, dtype=tf.int32, name='num_pos_anchor')

    valid_anchor_labels = tf.boolean_mask(anchor_labels, valid_mask)
  valid_label_logits = tf.boolean_mask(label_logits, valid_mask)

  with tf.name_scope('label_metrics'):
    valid_label_prob = tf.nn.sigmoid(valid_label_logits)
    summaries = []
    with tf.device('/cpu:0'):
      for th in [0.5, 0.2, 0.1]:
        valid_prediction = tf.cast(valid_label_prob > th, tf.int32)
        nr_pos_prediction = tf.reduce_sum(valid_prediction, name='num_pos_prediction')
        pos_prediction_corr = tf.count_nonzero(
          tf.logical_and(
            valid_label_prob > th,
            tf.equal(valid_prediction, valid_anchor_labels)),
          dtype=tf.int32)
        summaries.append(tf.truediv(
          pos_prediction_corr,
          nr_pos, name='recall_th{}'.format(th)))
        precision = tf.to_float(tf.truediv(pos_prediction_corr, nr_pos_prediction))
        precision = tf.where(tf.equal(nr_pos_prediction, 0), 0.0, precision, name='precision_th{}'.format(th))
        summaries.append(precision)
    # add_moving_summary(*summaries)

  label_loss = tf.nn.sigmoid_cross_entropy_with_logits(
    labels=tf.to_float(valid_anchor_labels), logits=valid_label_logits)
  label_loss = tf.reduce_mean(label_loss, name='label_loss')

  pos_anchor_boxes = tf.boolean_mask(anchor_boxes, pos_mask)
  pos_box_logits = tf.boolean_mask(box_logits, pos_mask)
  delta = 1.0 / 9
  box_loss = tf.losses.huber_loss(
    pos_anchor_boxes, pos_box_logits, delta=delta,
    reduction=tf.losses.Reduction.SUM) / delta
  box_loss = tf.div(
    box_loss,
    tf.cast(nr_valid, tf.float32), name='box_loss')

  # add_moving_summary(label_loss, box_loss, nr_valid, nr_pos)
  return label_loss, box_loss def fastrcnn_losses(labels, label_logits, fg_boxes, fg_box_logits):
  """
  Args:
      labels: n,
      label_logits: nxC
      fg_boxes: nfgx4, encoded
      fg_box_logits: nfgx(C-1)x4
  """
  label_loss = tf.nn.sparse_softmax_cross_entropy_with_logits(
    labels=labels, logits=label_logits)
  label_loss = tf.reduce_mean(label_loss, name='label_loss')

  fg_inds = tf.where(labels > 0)[:, 0]
  fg_labels = tf.gather(labels, fg_inds)
  num_fg = tf.size(fg_inds)
  indices = tf.stack(
    [tf.range(num_fg),
     tf.to_int32(fg_labels) - 1], axis=1)  # #fgx2
  fg_box_logits = tf.gather_nd(fg_box_logits, indices)

  # with tf.name_scope('label_metrics'), tf.device('/cpu:0'):
  #   prediction = tf.argmax(label_logits, axis=1, name='label_prediction')
  #   correct = tf.to_float(tf.equal(prediction, labels))  # boolean/integer gather is unavailable on GPU
  #   accuracy = tf.reduce_mean(correct, name='accuracy')
  #   fg_label_pred = tf.argmax(tf.gather(label_logits, fg_inds), axis=1)
  #   num_zero = tf.reduce_sum(tf.to_int32(tf.equal(fg_label_pred, 0)), name='num_zero')
  #   false_negative = tf.truediv(num_zero, num_fg, name='false_negative')
  #   fg_accuracy = tf.reduce_mean(
  #     tf.gather(correct, fg_inds), name='fg_accuracy')

  box_loss = tf.losses.huber_loss(
    fg_boxes, fg_box_logits, reduction=tf.losses.Reduction.SUM)
  box_loss = tf.truediv(
    box_loss, tf.to_float(tf.shape(labels)[0]), name='box_loss')

  # add_moving_summary(label_loss, box_loss, accuracy, fg_accuracy, false_negative)
  return label_loss, box_loss def clip_boxes(boxes, window, name=None):
  """
  Args:
      boxes: nx4, xyxy
      window: [h, w]
  """
  boxes = tf.maximum(boxes, 0.0)
  m = tf.tile(tf.reverse(window, [0]), [2])  # (4,)
  boxes = tf.minimum(boxes, tf.to_float(m), name=name)
  return boxes def fastrcnn_predictions(boxes, probs, num_classes, result_score_thresh, class_agnostic_boxes=False):
  """
  Generate final results from predictions of all proposals.
  Args:
      boxes: n#catx4 floatbox in float32
      probs: nx#class
  """
  if class_agnostic_boxes:
    assert boxes.shape[1] == 1, (boxes.shape[1], 1)
  else:
    assert boxes.shape[1] == num_classes - 1, (boxes.shape[1], num_classes - 1)
  assert probs.shape[1] == num_classes, (probs.shape[1], num_classes)
  boxes = tf.transpose(boxes, [1, 0, 2])  # #catxnx4
  probs = tf.transpose(probs[:, 1:], [1, 0])  # #catxn

  if class_agnostic_boxes:
    num_cat = tf.shape(probs)[0]
    boxes = tf.tile(boxes, [num_cat, 1, 1])

  def f(X):
    """
    prob: n probabilities
    box: nx4 boxes
    Returns: n boolean, the selection
    """
    prob, box = X
    output_shape = tf.shape(prob)
    # filter by score threshold
    ids = tf.reshape(tf.where(prob > result_score_thresh), [-1])
    prob = tf.gather(prob, ids)
    box = tf.gather(box, ids)
    # NMS within each class
    # Version from tensorflow 1.10, where everything works
    from tensorflow.python.ops import gen_image_ops
    tf.image.non_max_suppression = gen_image_ops.non_max_suppression  # or v3, v4
    selection = tf.image.non_max_suppression(
      box, prob, RESULTS_PER_IM, FASTRCNN_NMS_THRESH)
    selection = tf.to_int32(tf.gather(ids, selection))
    # sort available in TF>1.4.0
    # sorted_selection = tf.contrib.framework.sort(selection, direction='ASCENDING')
    sorted_selection = -tf.nn.top_k(-selection, k=tf.size(selection))[0]
    mask = tf.sparse_to_dense(
      sparse_indices=sorted_selection,
      output_shape=output_shape,
      sparse_values=True,
      default_value=False)
    return mask

  masks = tf.map_fn(f, (probs, boxes), dtype=tf.bool,
                    parallel_iterations=10)  # #cat x N
  selected_indices = tf.where(masks)  # #selection x 2, each is (cat_id, box_id)
  probs = tf.boolean_mask(probs, masks)

  # filter again by sorting scores
  topk_probs, topk_indices = tf.nn.top_k(
    probs,
    tf.minimum(RESULTS_PER_IM, tf.size(probs)),
    sorted=False)
  filtered_selection = tf.gather(selected_indices, topk_indices)
  filtered_selection = tf.reverse(filtered_selection, axis=[1], name='filtered_indices')
  return filtered_selection, topk_probs def maskrcnn_loss(mask_logits, fg_labels, fg_target_masks, class_agnostic):
  """
  Args:
      mask_logits: #fg x #category x14x14
      fg_labels: #fg, in 1~#class
      fg_target_masks: #fgx14x14, int
  """
  if class_agnostic:
    # probably we could also set this to tf.ones?
    fg_labels = tf.cast(tf.greater(fg_labels, 0), fg_labels.dtype)
  num_fg = tf.size(fg_labels)
  indices = tf.stack([tf.range(num_fg), tf.to_int32(fg_labels) - 1], axis=1)  # #fgx2
  # mask_logits needs to be NCHW for the next operation
  mask_logits = tf.transpose(mask_logits, [0, 3, 1, 2])
  mask_logits = tf.gather_nd(mask_logits, indices)  # #fgx14x14
  mask_probs = tf.sigmoid(mask_logits)

  # add some training visualizations to tensorboard
  with tf.name_scope('mask_viz'):
    viz = tf.concat([fg_target_masks, mask_probs], axis=1)
    viz = tf.expand_dims(viz, 3)
    viz = tf.cast(viz * 255, tf.uint8, name='viz')
    tf.summary.image('mask_truth|pred', viz, max_outputs=10)

  loss = tf.nn.sigmoid_cross_entropy_with_logits(
    labels=fg_target_masks, logits=mask_logits)
  loss = tf.reduce_mean(loss, name='maskrcnn_loss')

  pred_label = mask_probs > 0.5
  truth_label = fg_target_masks > 0.5
  accuracy = tf.reduce_mean(
    tf.to_float(tf.equal(pred_label, truth_label)),
    name='accuracy')
  pos_accuracy = tf.logical_and(
    tf.equal(pred_label, truth_label),
    tf.equal(truth_label, True))
  pos_accuracy = tf.reduce_mean(tf.to_float(pos_accuracy), name='pos_accuracy')
  fg_pixel_ratio = tf.reduce_mean(tf.to_float(truth_label), name='fg_pixel_ratio')

  # add_moving_summary(loss, accuracy, fg_pixel_ratio, pos_accuracy)
  return loss def crop_and_resize(image, boxes, box_ind, crop_size, pad_border=True):
  """
  Better-aligned version of tf.image.crop_and_resize, following our definition of floating point boxes.
  Args:
      image: NHWC
      boxes: nx4, x1y1x2y2
      box_ind: (n,)
      crop_size (int):
  Returns:
      n,size,size,C
  """
  assert isinstance(crop_size, int), crop_size

  # TF's crop_and_resize produces zeros on border
  if pad_border:
    # this can be quite slow
    image = tf.pad(image, [[0, 0], [1, 1], [1, 1], [0, 0]], mode='SYMMETRIC')
    boxes = boxes + 1

  def transform_fpcoor_for_tf(boxes, image_shape, crop_shape):
    """
    The way tf.image.crop_and_resize works (with normalized box):
    Initial point (the value of output[0]): x0_box * (W_img - 1)
    Spacing: w_box * (W_img - 1) / (W_crop - 1)
    Use the above grid to bilinear sample.
    However, what we want is (with fpcoor box):
    Spacing: w_box / W_crop
    Initial point: x0_box + spacing/2 - 0.5
    (-0.5 because bilinear sample assumes floating point coordinate (0.0, 0.0) is the same as pixel value (0, 0))
    This function transform fpcoor boxes to a format to be used by tf.image.crop_and_resize
    Returns:
        y1x1y2x2
    """
    x0, y0, x1, y1 = tf.split(boxes, 4, axis=1)

    spacing_w = (x1 - x0) / tf.to_float(crop_shape[1])
    spacing_h = (y1 - y0) / tf.to_float(crop_shape[0])

    nx0 = (x0 + spacing_w / 2 - 0.5) / tf.to_float(image_shape[1] - 1)
    ny0 = (y0 + spacing_h / 2 - 0.5) / tf.to_float(image_shape[0] - 1)

    nw = spacing_w * tf.to_float(crop_shape[1] - 1) / tf.to_float(image_shape[1] - 1)
    nh = spacing_h * tf.to_float(crop_shape[0] - 1) / tf.to_float(image_shape[0] - 1)

    return tf.concat([ny0, nx0, ny0 + nh, nx0 + nw], axis=1)

  image_shape = tf.shape(image)[1:3]
  boxes = transform_fpcoor_for_tf(boxes, image_shape, [crop_size, crop_size])
  ret = tf.image.crop_and_resize(
    image, boxes, box_ind,
    crop_size=[crop_size, crop_size])
  return ret ALL_ANCHORS = None NUM_ANCHOR = 15  # This is hard coded ANCHOR_STRIDE = 16
import tensorflow as tf
import numpy as np
      import cv2
from functools import partial
from datasets import DataKeys
from network.ConvolutionalLayers import Conv, ConvTranspose
from network.FullyConnected import FullyConnected
from network.Layer import Layer
from network.Resnet import add_resnet_conv5
from core import Measures, Extractions
from network.FasterRCNN_utils import decode_bbox_target, encode_bbox_target,\
  generate_rpn_proposals, sample_fast_rcnn_targets, roi_align, rpn_losses,\
  fastrcnn_losses, clip_boxes, fastrcnn_predictions, maskrcnn_loss, crop_and_resize
from datasets.util.Detection import ALL_ANCHORS, NUM_ANCHOR, ANCHOR_STRIDE
    else:
      assert len(fastrcnn_reid_features) == 1
      fastrcnn_reid_features = fastrcnn_reid_features[0]
    return fastrcnn_box_logits, fastrcnn_label_logits, feature_fastrcnn, fastrcnn_reid_features

  def _create_losses(self, featuremap_labels, featuremap_boxes, fm_anchors, fm_shape, fastrcnn_box_logits,
                     fastrcnn_label_logits, feature_fastrcnn, fg_inds_wrt_gt, gt_boxes, gt_masks,
                     rcnn_labels, rcnn_sampled_boxes, rpn_box_logits, rpn_label_logits):
    anchor_labels = tf.slice(
      featuremap_labels, [0, 0, 0],
      tf.stack([fm_shape[0], fm_shape[1], -1]),
      name='sliced_anchor_labels')
    anchor_boxes = tf.slice(
      featuremap_boxes, [0, 0, 0, 0],
      tf.stack([fm_shape[0], fm_shape[1], -1, -1]),
      name='sliced_anchor_boxes')
    anchor_boxes_encoded = encode_bbox_target(anchor_boxes, fm_anchors)

    # rpn loss
    rpn_label_loss, rpn_box_loss = rpn_losses(
      anchor_labels, anchor_boxes_encoded, rpn_label_logits, rpn_box_logits)
    # fastrcnn loss
    fg_inds_wrt_sample = tf.reshape(tf.where(rcnn_labels > 0), [-1])  # fg inds w.r.t all samples
    fg_sampled_boxes = tf.gather(rcnn_sampled_boxes, fg_inds_wrt_sample)
    matched_gt_boxes = tf.gather(gt_boxes, fg_inds_wrt_gt)
    encoded_boxes = encode_bbox_target(
      matched_gt_boxes,
      fg_sampled_boxes) * tf.constant(FASTRCNN_BBOX_REG_WEIGHTS)
    fastrcnn_label_loss, fastrcnn_box_loss = fastrcnn_losses(
      rcnn_labels, fastrcnn_label_logits,
      encoded_boxes,
      tf.gather(fastrcnn_box_logits, fg_inds_wrt_sample))

    if self._add_maskrcnn:
      # maskrcnn loss
      fg_labels = tf.gather(rcnn_labels, fg_inds_wrt_sample)
      fg_feature = tf.gather(feature_fastrcnn, fg_inds_wrt_sample)
      mask_logits = maskrcnn_head(fg_feature, self._num_classes, self.tower_setup,
                                  class_agnostic_conv=self.class_agnostic_box_and_mask_heads)  # #fg x #cat x 14x14

      gt_masks_for_fg = tf.gather(gt_masks, fg_inds_wrt_gt)  # nfg x H x W
      target_masks_for_fg = crop_and_resize(
        tf.expand_dims(gt_masks_for_fg, 3),
        fg_sampled_boxes,
        tf.range(tf.size(fg_inds_wrt_gt)), 14, pad_border=False)  # nfg x 1x14x14
      target_masks_for_fg = tf.squeeze(target_masks_for_fg, 3, 'sampled_fg_mask_targets')
      mrcnn_loss = maskrcnn_loss(mask_logits, fg_labels, target_masks_for_fg,
                                 class_agnostic=self.class_agnostic_box_and_mask_heads)
      self.add_scalar_summary(mrcnn_loss, "mrcnn_loss")
    else:
      mrcnn_loss = 0.0

    return fastrcnn_box_loss, fastrcnn_label_loss, mrcnn_loss, rpn_box_loss, rpn_label_loss

  def _create_final_outputs(self, batch_idx, rpn_input, fastrcnn_box_logits, fastrcnn_label_logits,
                            proposal_boxes, reid_features):
    label_probs = tf.nn.softmax(fastrcnn_label_logits, name='fastrcnn_all_probs')  # #proposal x #Class
    if self.class_agnostic_box_and_mask_heads:
      anchors = tf.tile(tf.expand_dims(proposal_boxes, 1), [1, 1, 1])  # #proposal x #Cat x 4
    else:
      anchors = tf.tile(tf.expand_dims(proposal_boxes, 1), [1, self._num_classes - 1, 1])  # #proposal x #Cat x 4
    decoded_boxes = decode_bbox_target(
      self.bbox_decode_clip,
      fastrcnn_box_logits /
      tf.constant(FASTRCNN_BBOX_REG_WEIGHTS), anchors)
    decoded_boxes = clip_boxes(decoded_boxes, self._image_shape2d, name='fastrcnn_all_boxes')
    if self.provide_boxes_as_input:
      n_proposals = tf.shape(proposal_boxes)[0]
      pred_indices = tf.stack([tf.range(n_proposals), tf.zeros((n_proposals,), dtype=tf.int32)], axis=1)
      final_probs = tf.ones((n_proposals,))
      ## do not do a second bbox regression!
      #decoded_boxes = anchors
    else:
      # indices: Nx2. Each index into (#proposal, #category)
      pred_indices, final_probs = fastrcnn_predictions(decoded_boxes, label_probs,
                                                       self._num_classes, self.result_score_thresh,
                                                       self.class_agnostic_box_and_mask_heads)
    final_probs = tf.identity(final_probs, 'final_probs')
    if self.class_agnostic_box_and_mask_heads:
      final_boxes = tf.gather(tf.squeeze(decoded_boxes, axis=1), pred_indices[:, 0], name='final_boxes')
    else:
      final_boxes = tf.gather_nd(decoded_boxes, pred_indices, name='final_boxes')
    final_labels = tf.add(pred_indices[:, 1], 1, name='final_labels')
    self.outputs_per_batch_idx[batch_idx] = [final_probs, final_boxes, final_labels]

    if self._add_maskrcnn:
      # HACK to work around https://github.com/tensorflow/tensorflow/issues/14657
      def f1():
        roi_resized = roi_align(rpn_input, final_boxes * (1.0 / ANCHOR_STRIDE), 14)
        with tf.variable_scope(tf.get_variable_scope(), reuse=True):
          feature_maskrcnn = add_resnet_conv5(roi_resized, self.tower_setup)[0]
        mask_logits = maskrcnn_head(
          feature_maskrcnn, self._num_classes, self.tower_setup,
          class_agnostic_conv=self.class_agnostic_box_and_mask_heads)  # #result x #cat x 14x14
        mask_logits = tf.transpose(mask_logits, [0, 3, 1, 2])
        if self.class_agnostic_box_and_mask_heads:
          mask_logits = tf.squeeze(mask_logits, axis=1)
        else:
          indices = tf.stack([tf.range(tf.size(final_labels)), tf.to_int32(final_labels) - 1], axis=1)
          mask_logits = tf.gather_nd(mask_logits, indices)  # #resultx14x14
        return tf.sigmoid(mask_logits)

      final_masks = tf.cond(tf.size(final_probs) > 0, f1, lambda: tf.zeros([0, 14, 14]))
      final_masks = tf.identity(final_masks, name='final_masks')
      self.outputs_per_batch_idx[batch_idx].append(final_masks)
    else:
      final_masks = None

    if self._do_reid:
      final_reid_features = tf.gather(reid_features, pred_indices[:, 0])
      self.outputs_per_batch_idx[batch_idx].append(final_reid_features)
    else:
      final_reid_features = None

    return final_boxes, final_labels, final_masks, final_probs, final_reid_features

  def _add_basic_measures(self, inp, loss):
    n_examples = tf.shape(inp)[0]
    self.measures[Measures.N_EXAMPLES] = n_examples
    if loss is not None:

      self.measures[Measures.LOSS] = loss * tf.cast(n_examples, tf.float32)