jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def create_dataset_from_data_producer(producer, params):
  """Return dataset online-generating data."""

  def preprocess_train_input(features, labels):
    """Pre-process the training data.

    This is needed because
    - The label needs to be extended to be used in the loss fn
    - We need the same inputs for training and eval so adding fake inputs
      for DUPLICATE_MASK in training data.

    Args:
      features: Dictionary of features for training.
      labels: Training labels.

    Returns:
      Processed training features.
    """
    fake_dup_mask = tf.zeros_like(features[movielens.USER_COLUMN])
    features[rconst.DUPLICATE_MASK] = fake_dup_mask
    features[rconst.TRAIN_LABEL_KEY] = labels
    return features

  train_input_fn = producer.make_input_fn(is_training=True)
  train_input_dataset = train_input_fn(params).map(preprocess_train_input)

  def preprocess_eval_input(features):
    """Pre-process the eval data.

    This is needed because:
    - The label needs to be extended to be used in the loss fn
    - We need the same inputs for training and eval so adding fake inputs
      for VALID_PT_MASK in eval data.

    Args:
      features: Dictionary of features for evaluation.

    Returns:
      Processed evaluation features.
    """
    labels = tf.cast(tf.zeros_like(features[movielens.USER_COLUMN]), tf.bool)
    fake_valid_pt_mask = tf.cast(
        tf.zeros_like(features[movielens.USER_COLUMN]), tf.bool)
    features[rconst.VALID_POINT_MASK] = fake_valid_pt_mask
    features[rconst.TRAIN_LABEL_KEY] = labels
    return features

  eval_input_fn = producer.make_input_fn(is_training=False)
  eval_input_dataset = eval_input_fn(params).map(preprocess_eval_input)

  return train_input_dataset, eval_input_dataset
//  [python]
def _compute_accuracy(self, labels, logits):
        labels = labels[:, self.config.n_loss_skip:]
        pred = tf.argmax(logits, -1, output_type=labels.dtype)
        pred = pred[:, self.config.n_loss_skip:]
        tf.debugging.assert_equal(tf.shape(pred), tf.shape(labels))
        return tf.to_float(pred == labels)
//  [python]
def _compute_accuracy(self, labels, logits):
        labels = labels[:, self.config.n_loss_skip:]
        pred = tf.argmax(logits, -1, output_type=labels.dtype)
        pred = pred[:, self.config.n_loss_skip:]
        tf.debugging.assert_equal(tf.shape(pred), tf.shape(labels))
        return tf.cast(pred == labels, tf.float32)
//  [python]
def resize_and_pad_image(image, tar_shape = [896,1152], jitter=None):
    
    image_shape = tf.to_int32(tf.shape(image)[:2])
    ratio = tar_shape[0] / image_shape[0]
    if ratio*image_shape[1] > tar_shape[1]:
        ratio = tar_shape[1] / image_shape[1]
    image_shape = ratio * image_shape
    image = tf.image.resize(image, tf.cast(image_shape, dtype=tf.int32))
    image = tf.image.pad_to_bounding_box(
        image, 0, 0, tar_shape[0], tar_shape[1]
    )
    return image, image_shape, ratio
//  [python]
def resize_and_pad_image(image, tar_shape = [896,1152], jitter=None):
    
    image_shape = tf.cast(tf.shape(image)[:2], dtype=num_precision)
    ratio = tar_shape[0] / image_shape[0]
    if ratio*image_shape[1] > tar_shape[1]:
        ratio = tar_shape[1] / image_shape[1]
    image_shape = ratio * image_shape
    image = tf.image.resize(image, tf.cast(image_shape, dtype=tf.int32))
    image = tf.image.pad_to_bounding_box(
        image, 0, 0, tar_shape[0], tar_shape[1]
    )
    return image, image_shape, ratio
//  [python]
def resize_and_crop_image(image, tar_shape = [896,1280],  #[1024,1536], #
    jitter=None):
    
    image_shape = tf.to_float(tf.shape(image)[-3:-1])
    ratio = tar_shape[1] / image_shape[1]
    if ratio*image_shape[0] < tar_shape[0]:
        ratio = tar_shape[0] / image_shape[0]
    image_shape = ratio * image_shape
    image = tf.image.resize(image, tf.to_int32(image_shape))
    image = tf.image.crop_to_bounding_box(
        image, 0, 0, tar_shape[0], tar_shape[1]
    )
    return image, image_shape, ratio
//  [python]
def resize_and_crop_image(image, tar_shape = [896,1280],  #[1024,1536], #
    jitter=None):
    
    image_shape = tf.cast(tf.shape(image)[-3:-1], dtype=num_precision)
    ratio = tar_shape[1] / image_shape[1]
    if ratio*image_shape[0] < tar_shape[0]:
        ratio = tar_shape[0] / image_shape[0]
    image_shape = ratio * image_shape
    image = tf.image.resize(image, tf.cast(image_shape, dtype=tf.int32))
    image = tf.image.crop_to_bounding_box(
        image, 0, 0, tar_shape[0], tar_shape[1]
    )
    return image, image_shape, ratio
//  [python]
def call(self, sample, rand_flip=True, return_ratio=False):
        
        image = sample["image"]
        bbox = sample["objects"]["bbox"]
        class_id = t.cast(sample["objects"]["label"], dtype=t.int32)

        if self.rand_flip or rand_flip: image, bbox = random_flip_horizontal(image, bbox)
        image, image_shape, ratio = resize_and_crop_image(image, tar_shape=self.tar_shape)

        bbox = tf.stack(
            [
                bbox[:, 0] * image_shape[1],
                bbox[:, 1] * image_shape[0],
                bbox[:, 2] * image_shape[1],
                bbox[:, 3] * image_shape[0],
            ],
            axis=-1,
        )
        bbox = convert_to_xywh(bbox)
        if self.return_ratio or return_ratio: return image, bbox, class_id, ratio
        return image, bbox, class_id
//  [python]
def call(self, sample, rand_flip=True, return_ratio=False):
        
        image = sample["image"]
        bbox = sample["objects"]["bbox"]
        class_id = tf.cast(sample["objects"]["label"], dtype=tf.int32)

        if self.rand_flip or rand_flip: image, bbox = random_flip_horizontal(image, bbox)
        image, image_shape, ratio = resize_and_crop_image(image, tar_shape=self.tar_shape)

        bbox = tf.stack(
            [
                bbox[:, 0] * image_shape[1],
                bbox[:, 1] * image_shape[0],
                bbox[:, 2] * image_shape[1],
                bbox[:, 3] * image_shape[0],
            ],
            axis=-1,
        )
        bbox = convert_to_xywh(bbox)
        if self.return_ratio or return_ratio: return image, bbox, class_id, ratio
        return image, bbox, class_id
//  [python]
def preprocess_data(sample, return_ratio=False, rand_flip=True):
    """Applies preprocessing step to a single sample

    Arguments:
      sample: A dict representing a single training sample.

    Returns:
      image: Resized and padded image with random horizontal flipping applied.
      bbox: Bounding boxes with the shape `(num_objects, 4)` where each box is
        of the format `[x, y, width, height]`.
      class_id: An tensor representing the class id of the objects, having
        shape `(num_objects,)`.
    """
    image = sample["image"]
    bbox = swap_xy(sample["objects"]["bbox"])
    class_id = t.cast(sample["objects"]["label"], dtype=t.int32)

    if rand_flip: image, bbox = random_flip_horizontal(image, bbox)
    image, image_shape, ratio = resize_and_pad_image(image)

    bbox = tf.stack(
        [
            bbox[:, 0] * image_shape[1],
            bbox[:, 1] * image_shape[0],
            bbox[:, 2] * image_shape[1],
            bbox[:, 3] * image_shape[0],
        ],
        axis=-1,
    )
    bbox = convert_to_xywh(bbox)
    if return_ratio: return image, bbox, class_id, ratio
    return image, bbox, class_id
//  [python]
def preprocess_data(sample, return_ratio=False, rand_flip=True):
    """Applies preprocessing step to a single sample

    Arguments:
      sample: A dict representing a single training sample.

    Returns:
      image: Resized and padded image with random horizontal flipping applied.
      bbox: Bounding boxes with the shape `(num_objects, 4)` where each box is
        of the format `[x, y, width, height]`.
      class_id: An tensor representing the class id of the objects, having
        shape `(num_objects,)`.
    """
    image = sample["image"]
    bbox = swap_xy(sample["objects"]["bbox"])
    class_id = tf.cast(sample["objects"]["label"], dtype=tf.int32)

    if rand_flip: image, bbox = random_flip_horizontal(image, bbox)
    image, image_shape, ratio = resize_and_pad_image(image)

    bbox = tf.stack(
        [
            bbox[:, 0] * image_shape[1],
            bbox[:, 1] * image_shape[0],
            bbox[:, 2] * image_shape[1],
            bbox[:, 3] * image_shape[0],
        ],
        axis=-1,
    )
    bbox = convert_to_xywh(bbox)
    if return_ratio: return image, bbox, class_id, ratio
    return image, bbox, class_id
//  [python]
def call(self, images, predictions):
        predictions = self.adapt(predictions)
        image_shape = tf.to_int32(tf.shape(images))
        anchor_boxes = self._anchor_box.get_anchors(image_shape[1], image_shape[2])

        cls_predictions = tf.nn.sigmoid(predictions[:, :, 4:])
        boxes = self._decode_box_predictions(anchor_boxes[None, ...], predictions[:, :, :4])

        return tf.image.combined_non_max_suppression(
            tf.expand_dims(boxes, axis=2),
            cls_predictions,
            self.max_detections_per_class,
            self.max_detections,
            self.nms_iou_threshold,
            self.confidence_threshold,
            clip_boxes=False,
        )
//  [python]
def call(self, images, predictions):
        predictions = self.adapt(predictions)
        image_shape = tf.cast(tf.shape(images), dtype=num_precision)
        anchor_boxes = self._anchor_box.get_anchors(image_shape[1], image_shape[2])

        cls_predictions = tf.nn.sigmoid(predictions[:, :, 4:])
        boxes = self._decode_box_predictions(anchor_boxes[None, ...], predictions[:, :, :4])

        return tf.image.combined_non_max_suppression(
            tf.expand_dims(boxes, axis=2),
            cls_predictions,
            self.max_detections_per_class,
            self.max_detections,
            self.nms_iou_threshold,
            self.confidence_threshold,
            clip_boxes=False,
        )
//  [python]
def _encode_sample(self, image_shape, gt_boxes, cls_ids):
        """Creates box and classification targets for a single sample"""
        if self.use_empty_samples and tf.reduce_max(cls_ids)==0: 
            gt_boxes = tf.convert_to_tensor([[0.0, 0.0, 1.0, 1.0]], dtype=num_precision)
        anchor_boxes = self._anchor_box.get_anchors(image_shape[1], image_shape[2])
        cls_ids = tf.cast(cls_ids, dtype=num_precision)
        matched_gt_idx, positive_mask, ignore_mask = self._match_anchor_boxes(
            anchor_boxes, gt_boxes
        )
        matched_gt_boxes = tf.gather(gt_boxes, matched_gt_idx)
        box_target = self._compute_box_target(anchor_boxes, matched_gt_boxes)
        matched_gt_cls_ids = tf.gather(cls_ids, matched_gt_idx)
        cls_target = tf.where(
            tf.not_equal(positive_mask, 1.0), -1.0, matched_gt_cls_ids
        )
        cls_target = tf.where(tf.equal(ignore_mask, 1.0), -2.0, cls_target)
        cls_target = tf.expand_dims(cls_target, axis=-1)
        label = tf.concat([box_target, cls_target], axis=-1)
        return label
//  [python]
def call(self, input_data):
        out = tf.convert_to_tensor(input_data)
        out = t.cast(out, dtype=self.prec)
        if not self.norm_adapted:
            self.adapt_norm(tf.shape(out))
        out = out*(self.norm[0]-self.norm[1]) + self.norm[1]
        return t.cast(out, dtype=num_precision)
//  [python]
def call(self, input_data):
        out = tf.convert_to_tensor(input_data)
        out = tf.cast(out, dtype=self.prec)
        if not self.norm_adapted:
            self.adapt_norm(tf.shape(out))
        out = out*(self.norm[0]-self.norm[1]) + self.norm[1]
        return tf.cast(out, dtype=num_precision)
//  [python]
def call(self, y_true, y_pred):
        y_pred = t.cast(y_pred, dtype=num_precision)
        box_labels = y_true[:, :, :4]
        cls_labels = t.one_hot(
            t.cast(y_true[:, :, 4], dtype=t.int32),
            depth=self._num_classes,
            dtype=num_precision,
        )
        y_pred = self.adapt(y_pred)
        box_predictions = y_pred[:, :, :4]
        cls_predictions = y_pred[:, :, 4:]

        positive_mask = t.cast(t.greater(y_true[:, :, 4], -1.0), dtype=num_precision)
        ignore_mask = t.cast(t.equal(y_true[:, :, 4], -2.0), dtype=num_precision)
        clf_loss = self._clf_loss(cls_labels, cls_predictions)
        box_loss = self._box_loss(box_labels, box_predictions)
        clf_loss = t.where(t.equal(ignore_mask, 1.0), 0.0, clf_loss)
        box_loss = t.where(t.equal(positive_mask, 1.0), box_loss, 0.0)
        normalizer = t.reduce_sum(positive_mask, axis=-1)
        clf_loss = t.math.divide_no_nan(t.reduce_sum(clf_loss, axis=-1), normalizer)
        box_loss = t.math.divide_no_nan(t.reduce_sum(box_loss, axis=-1), normalizer)
        loss = clf_loss + box_loss
        return loss
//  [python]
def call(self, y_true, y_pred):
        y_pred = tf.cast(y_pred, dtype=num_precision)
        box_labels = y_true[:, :, :4]
        cls_labels = tf.one_hot(
            tf.cast(y_true[:, :, 4], dtype=tf.int32),
            depth=self._num_classes,
            dtype=num_precision,
        )
        y_pred = self.adapt(y_pred)
        box_predictions = y_pred[:, :, :4]
        cls_predictions = y_pred[:, :, 4:]

        positive_mask = tf.cast(tf.greater(y_true[:, :, 4], -1.0), dtype=num_precision)
        ignore_mask = tf.cast(tf.equal(y_true[:, :, 4], -2.0), dtype=num_precision)
        clf_loss = self._clf_loss(cls_labels, cls_predictions)
        box_loss = self._box_loss(box_labels, box_predictions)
        clf_loss = tf.where(tf.equal(ignore_mask, 1.0), 0.0, clf_loss)
        box_loss = tf.where(tf.equal(positive_mask, 1.0), box_loss, 0.0)
        normalizer = tf.reduce_sum(positive_mask, axis=-1)
        clf_loss = tf.math.divide_no_nan(tf.reduce_sum(clf_loss, axis=-1), normalizer)
        box_loss = tf.math.divide_no_nan(tf.reduce_sum(box_loss, axis=-1), normalizer)
        loss = clf_loss + box_loss
        return loss
//  [python]
def call(self, images, predictions):
        image_shape = tf.to_int32(tf.shape(images))
        anchor_boxes = self._anchor_box.get_anchors(image_shape[1], image_shape[2])
        ps = predictions.shape
        cls_predictions = np.zeros([1,ps[1],self.num_classes], dtype='float32')
        o_idx = tf.squeeze(tf.where(tf.greater(predictions[0, :, 4], -1.0))).numpy()
        if not hasattr(o_idx, '__len__'): o_idx = [o_idx]
        for idx in o_idx:
            cls_predictions[0,idx,int(predictions[0,idx,4])] = 1.0
        boxes = self._decode_box_predictions(anchor_boxes[None, ...], predictions[:, :, :4])

        return tf.image.combined_non_max_suppression(
            tf.expand_dims(boxes, axis=2),
            cls_predictions,
            self.max_detections_per_class,
            self.max_detections,
            self.nms_iou_threshold,
            self.confidence_threshold,
            clip_boxes=False,
        )
//  [python]
def call(self, images, predictions):
        image_shape = tf.cast(tf.shape(images), dtype=num_precision)
        anchor_boxes = self._anchor_box.get_anchors(image_shape[1], image_shape[2])
        ps = predictions.shape
        cls_predictions = np.zeros([1,ps[1],self.num_classes], dtype='float32')
        o_idx = tf.squeeze(tf.where(tf.greater(predictions[0, :, 4], -1.0))).numpy()
        if not hasattr(o_idx, '__len__'): o_idx = [o_idx]
        for idx in o_idx:
            cls_predictions[0,idx,int(predictions[0,idx,4])] = 1.0
        boxes = self._decode_box_predictions(anchor_boxes[None, ...], predictions[:, :, :4])

        return tf.image.combined_non_max_suppression(
            tf.expand_dims(boxes, axis=2),
            cls_predictions,
            self.max_detections_per_class,
            self.max_detections,
            self.nms_iou_threshold,
            self.confidence_threshold,
            clip_boxes=False,
        )
//  [python]
def view_sample(sample, showlabels=True):
    x = t.cast(sample["image"], dtype=t.float32)
    image_shape = t.cast(t.shape(x)[-3:-1], dtype=t.float32)
    #ybox = sample["objects"]["bbox"]
    #ybox = swap_xy(sample[1])
    ybox = convert_to_corners(sample[1])
    ycls = t.cast(sample["objects"]["label"], dtype=t.float32)
    ybox = t.stack( # [classes, x0, y0, x1, y1]
        [
            ybox[..., 0] * image_shape[1],
            ybox[..., 1] * image_shape[0],
            ybox[..., 2] * image_shape[1],
            ybox[..., 3] * image_shape[0],
        ],
        axis=-1)

    visualize_datasample(x,ybox,ycls, showlabels=showlabels)
//  [python]
def view_sample(sample, showlabels=True):
    x = tf.cast(sample["image"], dtype=tf.float32)
    image_shape = tf.cast(tf.shape(x)[-3:-1], dtype=tf.float32)
    #ybox = sample["objects"]["bbox"]
    #ybox = swap_xy(sample[1])
    ybox = convert_to_corners(sample[1])
    ycls = tf.cast(sample["objects"]["label"], dtype=tf.float32)
    ybox = tf.stack( # [classes, x0, y0, x1, y1]
        [
            ybox[..., 0] * image_shape[1],
            ybox[..., 1] * image_shape[0],
            ybox[..., 2] * image_shape[1],
            ybox[..., 3] * image_shape[0],
        ],
        axis=-1)

    visualize_datasample(x,ybox,ycls, showlabels=showlabels)
//  [python]
def view_sample2(sample, swapxy=False, showlabels=True):
    x = t.cast(sample["image"], dtype=t.float32)
    image_shape = t.cast(t.shape(x)[-3:-1], dtype=t.float32)
    ybox = sample["objects"]["bbox"]
    if swapxy: ybox = swap_xy(ybox)
    #ybox = convert_to_corners(sample[1])
    ycls = t.cast(sample["objects"]["label"], dtype=t.float32)
    ybox = t.stack( # [classes, x0, y0, x1, y1]
        [
            ybox[..., 0] * image_shape[1],
            ybox[..., 1] * image_shape[0],
            ybox[..., 2] * image_shape[1],
            ybox[..., 3] * image_shape[0],
        ],
        axis=-1)

    visualize_datasample(x,ybox,ycls, colors=colorlist, showlabels=showlabels)
//  [python]
def view_sample2(sample, swapxy=False, showlabels=True):
    x = tf.cast(sample["image"], dtype=tf.float32)
    image_shape = tf.cast(tf.shape(x)[-3:-1], dtype=tf.float32)
    ybox = sample["objects"]["bbox"]
    if swapxy: ybox = swap_xy(ybox)
    #ybox = convert_to_corners(sample[1])
    ycls = tf.cast(sample["objects"]["label"], dtype=tf.float32)
    ybox = tf.stack( # [classes, x0, y0, x1, y1]
        [
            ybox[..., 0] * image_shape[1],
            ybox[..., 1] * image_shape[0],
            ybox[..., 2] * image_shape[1],
            ybox[..., 3] * image_shape[0],
        ],
        axis=-1)

    visualize_datasample(x,ybox,ycls, colors=colorlist, showlabels=showlabels)
//  [python]
def _repeat(x, n_repeats):
    rep = tf.transpose(tf.expand_dims(tf.ones(shape=tf.stack([n_repeats,])), 1), [1, 0])
    rep = tf.to_float(rep)
    x = tf.matmul(tf.reshape(x, (-1, 1)), rep)
    return tf.reshape(x, [-1])
//  [python]
def _repeat(x, n_repeats):
    rep = tf.transpose(tf.expand_dims(tf.ones(shape=tf.stack([n_repeats,])), 1), [1, 0])
    rep = tf.cast(rep, 'float32')
    x = tf.matmul(tf.reshape(x, (-1, 1)), rep)
    return tf.reshape(x, [-1])
//  [python]
def _voxel_hard_negative_classification_loss_unbatched(
    inputs_1, outputs_1, is_intermediate, gamma):
  """Loss function for input and outputs of batch size 1."""
  inputs_1, outputs_1 = _get_voxels_valid_inputs_outputs(
      inputs_1=inputs_1, outputs_1=outputs_1)
  if is_intermediate:
    logits = outputs_1[standard_fields.DetectionResultFields
                       .intermediate_object_semantic_voxels]
  else:
    logits = outputs_1[
        standard_fields.DetectionResultFields.object_semantic_voxels]
  labels = tf.reshape(
      inputs_1[standard_fields.InputDataFields.object_class_voxels], [-1])
  background_mask = tf.equal(labels, 0)
  num_background_points = tf.reduce_sum(
      tf.cast(background_mask, dtype=tf.int32))

  def loss_fn():
    """Loss function."""
    num_classes = logits.get_shape().as_list()[-1]
    if num_classes is None:
      raise ValueError('Number of classes is unknown.')
    masked_logits = tf.boolean_mask(logits, background_mask)
    masked_weights = tf.pow(
        1.0 - tf.reshape(tf.nn.softmax(masked_logits)[:, 0], [-1, 1]), gamma)
    num_points = tf.shape(masked_logits)[0]
    masked_weights = masked_weights * tf.cast(
        num_points, dtype=tf.float32) / tf.reduce_sum(masked_weights)
    masked_labels_one_hot = tf.one_hot(
        indices=tf.boolean_mask(labels, background_mask), depth=num_classes)
    loss = classification_loss_fn(
        logits=masked_logits,
        labels=masked_labels_one_hot,
        weights=masked_weights)
    return loss

  cond = tf.logical_and(
      tf.greater(num_background_points, 0), tf.greater(tf.shape(labels)[0], 0))
  return tf.cond(cond, loss_fn, lambda: tf.constant(0.0, dtype=tf.float32))
//  [python]
def _generate_multilevel_boxes(self) -> Dict[str, tf.Tensor]:
    """Generates multi-scale anchor boxes.

    Returns:
      An OrderedDict from level to anchor boxes of shape [height_l, width_l,
      num_anchors_per_location * 4].
    """
    multilevel_boxes = collections.OrderedDict()
    for level in range(self.min_level, self.max_level + 1):
      boxes_l = []
      feat_size_y = math.ceil(self.image_size[0] / 2**level)
      feat_size_x = math.ceil(self.image_size[1] / 2**level)
      stride_y = t.float32(self.image_size[0] / feat_size_y)
      stride_x = t.float32(self.image_size[1] / feat_size_x)
      x = tf.range(stride_x / 2, self.image_size[1], stride_x)
      y = tf.range(stride_y / 2, self.image_size[0], stride_y)
      xv, yv = tf.meshgrid(x, y)
      for scale in range(self.num_scales):
        for aspect_ratio in self.aspect_ratios:
          intermidate_scale = 2 ** (scale / self.num_scales)
          base_anchor_size = self.anchor_size * 2**level * intermidate_scale
          aspect_x = aspect_ratio**0.5
          aspect_y = aspect_ratio**-0.5
          half_anchor_size_x = base_anchor_size * aspect_x / 2.0
          half_anchor_size_y = base_anchor_size * aspect_y / 2.0
          # Tensor shape Nx4.
          boxes = tf.stack(
              [
                  yv - half_anchor_size_y,
                  xv - half_anchor_size_x,
                  yv + half_anchor_size_y,
                  xv + half_anchor_size_x,
              ],
              axis=-1,
          )
          boxes_l.append(boxes)
      # Concat anchors on the same level to tensor shape HxWx(Ax4).
      boxes_l = tf.concat(boxes_l, axis=-1)
      multilevel_boxes[str(level)] = boxes_l
    return multilevel_boxes
//  [python]
def _generate_multilevel_boxes(self) -> Dict[str, tf.Tensor]:
    """Generates multi-scale anchor boxes.

    Returns:
      An OrderedDict from level to anchor boxes of shape [height_l, width_l,
      num_anchors_per_location * 4].
    """
    multilevel_boxes = collections.OrderedDict()
    for level in range(self.min_level, self.max_level + 1):
      boxes_l = []
      feat_size_y = math.ceil(self.image_size[0] / 2**level)
      feat_size_x = math.ceil(self.image_size[1] / 2**level)
      stride_y = tf.cast(self.image_size[0] / feat_size_y, tf.float32)
      stride_x = tf.cast(self.image_size[1] / feat_size_x, tf.float32)
      x = tf.range(stride_x / 2, self.image_size[1], stride_x)
      y = tf.range(stride_y / 2, self.image_size[0], stride_y)
      xv, yv = tf.meshgrid(x, y)
      for scale in range(self.num_scales):
        for aspect_ratio in self.aspect_ratios:
          intermidate_scale = 2 ** (scale / self.num_scales)
          base_anchor_size = self.anchor_size * 2**level * intermidate_scale
          aspect_x = aspect_ratio**0.5
          aspect_y = aspect_ratio**-0.5
          half_anchor_size_x = base_anchor_size * aspect_x / 2.0
          half_anchor_size_y = base_anchor_size * aspect_y / 2.0
          # Tensor shape Nx4.
          boxes = tf.stack(
              [
                  yv - half_anchor_size_y,
                  xv - half_anchor_size_x,
                  yv + half_anchor_size_y,
                  xv + half_anchor_size_x,
              ],
              axis=-1,
          )
          boxes_l.append(boxes)
      # Concat anchors on the same level to tensor shape HxWx(Ax4).
      boxes_l = tf.concat(boxes_l, axis=-1)
      multilevel_boxes[str(level)] = boxes_l
    return multilevel_boxes
//  [python]
def call(
      self,
      inputs: tf.Tensor,
      states: Optional[States] = None,
      output_states: bool = True,
  ) -> Union[tf.Tensor, Tuple[tf.Tensor, States]]:
    """Calls the layer with the given inputs.

    Args:
      inputs: An input `tf.Tensor`.
      states: A `dict` of states such that, if any of the keys match for this
        layer, will overwrite the contents of the buffer(s). Expected keys
        include `state_prefix + '/pos_enc_frame_count'`.
      output_states: A `bool`. If True, returns the output tensor and output
        states. Returns just the output tensor otherwise.

    Returns:
      An output `tf.Tensor` (and optionally the states if `output_states=True`).

    Raises:
      ValueError: If using 'channels_first' data format.
    """
    states = dict(states) if states is not None else {}

    # Keep a count of frames encountered across input iterations in
    # num_frames to be able to accurately update the positional encoding.
    num_frames = tf.shape(inputs)[1]
    frame_count = tf.cast(states.get(self._frame_count_name, [0]), tf.int32)
    states[self._frame_count_name] = frame_count + num_frames

    if self._cache_encoding:
      pos_encoding = self._pos_encoding
    else:
      pos_encoding = self._get_pos_encoding(
          tf.shape(inputs), frame_count=frame_count)
    pos_encoding = tf.cast(pos_encoding, inputs.dtype)
    pos_encoding = self._rezero(pos_encoding)
    outputs = inputs + pos_encoding

    return (outputs, states) if output_states else outputs
//  [python]
def get_relative_pose(target_pose, source_pose, is_vec=True):
    """[summary]
    get relative pose from two absolution poses
    Arguments:
        target_pose_vec {[batch_size, 6]} -- [rz, ry, rx, tx, ty, tz],
            or [batch_size, 4, 4] if is_vec=False
        source_pose_vec {[batch_size, 6]} -- [rz, ry, rx, tx, ty, tz]
            or [batch_size, 4, 4] if is_vec=False
        both are abosulte (up to scale) pose w.r.t. the world coord
        Note: this is different from the original SfMLearner which is [tx, ty, tz, rx, ry, rz].
    Returns:
        transformation from target to source
    """
    # target_pose = tf_render.Print(target_pose, [target_pose, target_pose.shape], message="target_pose", summarize=16)
    # source_pose = tf_render.Print(source_pose, [source_pose, source_pose.shape], message="source_pose", summarize=16)
    if is_vec:
        target_pose = pose_vec2mat(target_pose, clip=False)
        source_pose = pose_vec2mat(source_pose, clip=False)

    #target_pose_inv = tf_render.linalg.inv(target_pose)
    target_pose_inv = get_ext_inv(target_pose, False)
    transform_mat = tf.matmul(source_pose, target_pose_inv)

    transform_mat = t.cast(transform_mat, dtype=tf.float32)
    return transform_mat
//  [python]
def get_relative_pose(target_pose, source_pose, is_vec=True):
    """[summary]
    get relative pose from two absolution poses
    Arguments:
        target_pose_vec {[batch_size, 6]} -- [rz, ry, rx, tx, ty, tz],
            or [batch_size, 4, 4] if is_vec=False
        source_pose_vec {[batch_size, 6]} -- [rz, ry, rx, tx, ty, tz]
            or [batch_size, 4, 4] if is_vec=False
        both are abosulte (up to scale) pose w.r.t. the world coord
        Note: this is different from the original SfMLearner which is [tx, ty, tz, rx, ry, rz].
    Returns:
        transformation from target to source
    """
    # target_pose = tf_render.Print(target_pose, [target_pose, target_pose.shape], message="target_pose", summarize=16)
    # source_pose = tf_render.Print(source_pose, [source_pose, source_pose.shape], message="source_pose", summarize=16)
    if is_vec:
        target_pose = pose_vec2mat(target_pose, clip=False)
        source_pose = pose_vec2mat(source_pose, clip=False)

    #target_pose_inv = tf_render.linalg.inv(target_pose)
    target_pose_inv = get_ext_inv(target_pose, False)
    transform_mat = tf.matmul(source_pose, target_pose_inv)

    transform_mat = tf.cast(transform_mat, dtype=tf.float32)
    return transform_mat
//  [python]
def classification_loss_using_mask_iou_func_unbatched(
    embeddings, instance_ids, sampled_embeddings,
    sampled_instance_ids, sampled_class_labels, sampled_logits,
    similarity_strategy, is_balanced):
  """Classification loss using mask iou.

  Args:
    embeddings: A tf.float32 tensor of size [n, f].
    instance_ids: A tf.int32 tensor of size [n].
    sampled_embeddings: A tf.float32 tensor of size [num_samples, f].
    sampled_instance_ids: A tf.int32 tensor of size [num_samples].
    sampled_class_labels: A tf.int32 tensor of size [num_samples, 1].
    sampled_logits: A tf.float32 tensor of size [num_samples, num_classes].
    similarity_strategy: Defines the method for computing similarity between
                         embedding vectors. Possible values are 'dotproduct' and
                         'distance'.
    is_balanced: If True, the per-voxel losses are re-weighted to have equal
      total weight for foreground vs. background voxels.

  Returns:
    A tf.float32 loss scalar tensor.
  """
  predicted_soft_masks = metric_learning_utils.embedding_centers_to_soft_masks(
      embedding=embeddings,
      centers=sampled_embeddings,
      similarity_strategy=similarity_strategy)
  predicted_masks = tf.to_float(
      tf.greater(predicted_soft_masks, 0.5))
  gt_masks = tf.to_float(
      tf.equal(
          tf.expand_dims(sampled_instance_ids, axis=1),
          tf.expand_dims(instance_ids, axis=0)))
  pairwise_iou = instance_segmentation_utils.points_mask_pairwise_iou(
      masks1=predicted_masks, masks2=gt_masks)
  num_classes = sampled_logits.get_shape().as_list()[1]
  sampled_class_labels_one_hot = tf.one_hot(
      indices=tf.reshape(sampled_class_labels, [-1]), depth=num_classes)
  sampled_class_labels_one_hot_fg = sampled_class_labels_one_hot[:, 1:]
  iou_coefs = tf.tile(tf.reshape(pairwise_iou, [-1, 1]), [1, num_classes - 1])
  sampled_class_labels_one_hot_fg *= iou_coefs
  sampled_class_labels_one_hot_bg = tf.maximum(1.0 - tf.math.reduce_sum(
      sampled_class_labels_one_hot_fg, axis=1, keepdims=True), 0.0)
  sampled_class_labels_one_hot = tf.concat(
      [sampled_class_labels_one_hot_bg, sampled_class_labels_one_hot_fg],
      axis=1)
  params = {}
  if is_balanced:
    weights = loss_utils.get_balanced_loss_weights_multiclass(
        labels=tf.expand_dims(sampled_instance_ids, axis=1))
    params['weights'] = weights
  return classification_loss_fn(
      logits=sampled_logits, labels=sampled_class_labels_one_hot, **params)
//  [python]
def classification_loss_using_mask_iou_func_unbatched(
    embeddings, instance_ids, sampled_embeddings,
    sampled_instance_ids, sampled_class_labels, sampled_logits,
    similarity_strategy, is_balanced):
  """Classification loss using mask iou.

  Args:
    embeddings: A tf.float32 tensor of size [n, f].
    instance_ids: A tf.int32 tensor of size [n].
    sampled_embeddings: A tf.float32 tensor of size [num_samples, f].
    sampled_instance_ids: A tf.int32 tensor of size [num_samples].
    sampled_class_labels: A tf.int32 tensor of size [num_samples, 1].
    sampled_logits: A tf.float32 tensor of size [num_samples, num_classes].
    similarity_strategy: Defines the method for computing similarity between
                         embedding vectors. Possible values are 'dotproduct' and
                         'distance'.
    is_balanced: If True, the per-voxel losses are re-weighted to have equal
      total weight for foreground vs. background voxels.

  Returns:
    A tf.float32 loss scalar tensor.
  """
  predicted_soft_masks = metric_learning_utils.embedding_centers_to_soft_masks(
      embedding=embeddings,
      centers=sampled_embeddings,
      similarity_strategy=similarity_strategy)
  predicted_masks = tf.cast(
      tf.greater(predicted_soft_masks, 0.5), dtype=tf.float32)
  gt_masks = tf.cast(
      tf.equal(
          tf.expand_dims(sampled_instance_ids, axis=1),
          tf.expand_dims(instance_ids, axis=0)),
      dtype=tf.float32)
  pairwise_iou = instance_segmentation_utils.points_mask_pairwise_iou(
      masks1=predicted_masks, masks2=gt_masks)
  num_classes = sampled_logits.get_shape().as_list()[1]
  sampled_class_labels_one_hot = tf.one_hot(
      indices=tf.reshape(sampled_class_labels, [-1]), depth=num_classes)
  sampled_class_labels_one_hot_fg = sampled_class_labels_one_hot[:, 1:]
  iou_coefs = tf.tile(tf.reshape(pairwise_iou, [-1, 1]), [1, num_classes - 1])
  sampled_class_labels_one_hot_fg *= iou_coefs
  sampled_class_labels_one_hot_bg = tf.maximum(1.0 - tf.math.reduce_sum(
      sampled_class_labels_one_hot_fg, axis=1, keepdims=True), 0.0)
  sampled_class_labels_one_hot = tf.concat(
      [sampled_class_labels_one_hot_bg, sampled_class_labels_one_hot_fg],
      axis=1)
  params = {}
  if is_balanced:
    weights = loss_utils.get_balanced_loss_weights_multiclass(
        labels=tf.expand_dims(sampled_instance_ids, axis=1))
    params['weights'] = weights
  return classification_loss_fn(
      logits=sampled_logits, labels=sampled_class_labels_one_hot, **params)
//  [python]
def apply_bias_act(x, act='linear', alpha=None, gain=None, lrmul=1, bias_var='bias'):
    b = tf.get_variable(bias_var, shape=[x.shape[1]], initializer=tf.initializers.zeros()) * lrmul
    return fused_bias_act(x, b=tf.cast(b, x.dtype), act=act, alpha=alpha, gain=gain)
//  [python]
def dense_layer(x, fmaps, gain=1, use_wscale=True, lrmul=1, weight_var='weight', svd=False, factorized=False, lambda_mask=None, reduce_dims=None, sv_factors=None, svd_center=False, svd_config='S'):
    if len(x.shape) > 2:
        x = tf.reshape(x, [-1, np.prod([d.value for d in x.shape[1:]])])
    w = get_svd_weight([x.shape[1].value, fmaps], 
                       gain=gain, use_wscale=use_wscale, lrmul=lrmul, weight_var=weight_var, 
                       spatial=False, 
                       svd=svd, 
                       factorized=factorized, 
                       sv_factors=sv_factors, 
                       svd_center=svd_center, 
                       lambda_mask=lambda_mask, 
                       reduce_dims=reduce_dims,
                       svd_config=svd_config)
    w = tf.cast(w, x.dtype)
    return tf.matmul(x, w)
//  [python]
def conv2d_layer(x, fmaps, kernel, up=False, down=False, resample_kernel=None, gain=1, use_wscale=True, lrmul=1, weight_var='weight', svd=False, factorized=False, sv_factors=None, svd_center=False, adapt_func='training.networks_stylegan2.apply_identity', lambda_mask=None, reduce_dims=None, spatial=False, svd_config='S'):
    assert not (up and down)
    assert kernel >= 1 and kernel % 2 == 1
    x_prev = x
    w = get_svd_weight([kernel, kernel, x.shape[1].value, fmaps], 
                       gain=gain, use_wscale=use_wscale, lrmul=lrmul, weight_var=weight_var, 
                       spatial=spatial, 
                       svd=svd, 
                       factorized=factorized, 
                       sv_factors=sv_factors, 
                       svd_center=svd_center, 
                       lambda_mask=lambda_mask,
                       reduce_dims=reduce_dims,
                       svd_config=svd_config)
    if up:
        x = upsample_conv_2d(x, t.cast(w, x.dtype), data_format='NCHW', k=resample_kernel)
    elif down:
        x = conv_downsample_2d(x, t.cast(w, x.dtype), data_format='NCHW', k=resample_kernel)
    else:
        x = t.nn.conv2d(x, t.cast(w, x.dtype), data_format='NCHW', strides=[1,1,1,1], padding='SAME')
    x = call_func_by_name(func_name=adapt_func, x=x, x_prev=x)
    return x
//  [python]
def conv2d_layer(x, fmaps, kernel, up=False, down=False, resample_kernel=None, gain=1, use_wscale=True, lrmul=1, weight_var='weight', svd=False, factorized=False, sv_factors=None, svd_center=False, adapt_func='training.networks_stylegan2.apply_identity', lambda_mask=None, reduce_dims=None, spatial=False, svd_config='S'):
    assert not (up and down)
    assert kernel >= 1 and kernel % 2 == 1
    x_prev = x
    w = get_svd_weight([kernel, kernel, x.shape[1].value, fmaps], 
                       gain=gain, use_wscale=use_wscale, lrmul=lrmul, weight_var=weight_var, 
                       spatial=spatial, 
                       svd=svd, 
                       factorized=factorized, 
                       sv_factors=sv_factors, 
                       svd_center=svd_center, 
                       lambda_mask=lambda_mask,
                       reduce_dims=reduce_dims,
                       svd_config=svd_config)
    if up:
        x = upsample_conv_2d(x, tf.cast(w, x.dtype), data_format='NCHW', k=resample_kernel)
    elif down:
        x = conv_downsample_2d(x, tf.cast(w, x.dtype), data_format='NCHW', k=resample_kernel)
    else:
        x = tf.nn.conv2d(x, tf.cast(w, x.dtype), data_format='NCHW', strides=[1,1,1,1], padding='SAME')
    x = call_func_by_name(func_name=adapt_func, x=x, x_prev=x)
    return x
//  [python]
def minibatch_stddev_layer(x, group_size=4, num_new_features=1):
    group_size = tf.minimum(group_size, tf.shape(x)[0])     # Minibatch must be divisible by (or smaller than) group_size.
    s = x.shape                                             # [NCHW]  Input shape.
    y = tf.reshape(x, [group_size, -1, num_new_features, s[1]//num_new_features, s[2], s[3]])   # [GMncHW] Split minibatch into M groups of size G. Split channels into n channel groups c.
    y = tf.to_float(y)                              # [GMncHW] Cast to FP32.
    y -= tf.reduce_mean(y, axis=0, keepdims=True)           # [GMncHW] Subtract mean over group.
    y = tf.reduce_mean(tf.square(y), axis=0)                # [MncHW]  Calc variance over group.
    y = tf.sqrt(y + 1e-8)                                   # [MncHW]  Calc stddev over group.
    y = tf.reduce_mean(y, axis=[2,3,4], keepdims=True)      # [Mn111]  Take average over fmaps and pixels.
    y = tf.reduce_mean(y, axis=[2])                         # [Mn11] Split channels into c channel groups
    y = tf.cast(y, x.dtype)                                 # [Mn11]  Cast back to original data type.
    y = tf.tile(y, [group_size, 1, s[2], s[3]])             # [NnHW]  Replicate over group and pixels.
    return tf.concat([x, y], axis=1)
//  [python]
def minibatch_stddev_layer(x, group_size=4, num_new_features=1):
    group_size = tf.minimum(group_size, tf.shape(x)[0])     # Minibatch must be divisible by (or smaller than) group_size.
    s = x.shape                                             # [NCHW]  Input shape.
    y = tf.reshape(x, [group_size, -1, num_new_features, s[1]//num_new_features, s[2], s[3]])   # [GMncHW] Split minibatch into M groups of size G. Split channels into n channel groups c.
    y = tf.cast(y, tf.float32)                              # [GMncHW] Cast to FP32.
    y -= tf.reduce_mean(y, axis=0, keepdims=True)           # [GMncHW] Subtract mean over group.
    y = tf.reduce_mean(tf.square(y), axis=0)                # [MncHW]  Calc variance over group.
    y = tf.sqrt(y + 1e-8)                                   # [MncHW]  Calc stddev over group.
    y = tf.reduce_mean(y, axis=[2,3,4], keepdims=True)      # [Mn111]  Take average over fmaps and pixels.
    y = tf.reduce_mean(y, axis=[2])                         # [Mn11] Split channels into c channel groups
    y = tf.cast(y, x.dtype)                                 # [Mn11]  Cast back to original data type.
    y = tf.tile(y, [group_size, 1, s[2], s[3]])             # [NnHW]  Replicate over group and pixels.
    return tf.concat([x, y], axis=1)
//  [python]
def call(self, inputs):
        inputs, unbatched, _ = convert_to_ragged_batch(inputs)

        skip_masks = None
        if self.skip_list:
            skip_masks = self.StaticHashTable.lookup(inputs.flat_values)
        elif self.skip_fn:
            skip_masks = tf.map_fn(
                self.skip_fn, inputs.flat_values, fn_output_signature="bool"
            )
        elif self.skip_py_fn:

            def string_fn(token):
                return self.skip_py_fn(token.numpy().decode("utf-8"))

            def int_fn(token):
                return self.skip_py_fn(token.numpy())

            py_fn = string_fn if inputs.dtype == tf.string else int_fn

            skip_masks = tf.map_fn(
                lambda x: tf.py_function(py_fn, [x], "bool"),
                inputs.flat_values,
                fn_output_signature="bool",
            )

        positions = tf.ragged.range(inputs.row_lengths())

        if skip_masks is not None:
            skip_masks = tf.logical_not(skip_masks)
            skip_masks.set_shape([None])
            positions = tf.ragged.boolean_mask(
                positions, inputs.with_flat_values(skip_masks)
            )
        # Figure out how many we are going to select.
        token_counts = t.cast(positions.row_lengths(), "float32")
        num_to_select = tf.random.stateless_binomial(
            shape=tf.shape(token_counts),
            seed=self._generator.make_seeds()[:, 0],
            counts=token_counts,
            probs=self.rate,
        )
        if self.max_swaps is not None:
            num_to_select = tf.math.minimum(num_to_select, self.max_swaps)
        num_to_select = tf.math.minimum(
            num_to_select, t.cast(positions.row_lengths(), "int32")
        )
        num_to_select = t.cast(num_to_select, "int64")

        def _swap(x):
            positions, inputs, num_to_select = x
            for _ in range(num_to_select):
                index = tf.random.stateless_uniform(
                    shape=[2],
                    minval=0,
                    maxval=tf.size(positions),
                    dtype="int32",
                    seed=self._generator.make_seeds()[:, 0],
                )
                index1, index2 = positions[index[0]], positions[index[1]]
                # swap items at the sampled indices with each other
                inputs = tf.tensor_scatter_nd_update(
                    inputs,
                    [[index1], [index2]],
                    [inputs[index2], inputs[index1]],
                )
            return inputs

        swapped = tf.map_fn(
            _swap,
            (positions, inputs, num_to_select),
            fn_output_signature=tf.RaggedTensorSpec(
                ragged_rank=positions.ragged_rank - 1, dtype=inputs.dtype
            ),
        )
        swapped.flat_values.set_shape([None])

        if unbatched:
            swapped = tf.squeeze(swapped, axis=0)
        return swapped
//  [python]
def call(self, inputs):
        inputs, unbatched, _ = convert_to_ragged_batch(inputs)

        skip_masks = None
        if self.skip_list:
            skip_masks = self.StaticHashTable.lookup(inputs.flat_values)
        elif self.skip_fn:
            skip_masks = tf.map_fn(
                self.skip_fn, inputs.flat_values, fn_output_signature="bool"
            )
        elif self.skip_py_fn:

            def string_fn(token):
                return self.skip_py_fn(token.numpy().decode("utf-8"))

            def int_fn(token):
                return self.skip_py_fn(token.numpy())

            py_fn = string_fn if inputs.dtype == tf.string else int_fn

            skip_masks = tf.map_fn(
                lambda x: tf.py_function(py_fn, [x], "bool"),
                inputs.flat_values,
                fn_output_signature="bool",
            )

        positions = tf.ragged.range(inputs.row_lengths())

        if skip_masks is not None:
            skip_masks = tf.logical_not(skip_masks)
            skip_masks.set_shape([None])
            positions = tf.ragged.boolean_mask(
                positions, inputs.with_flat_values(skip_masks)
            )
        # Figure out how many we are going to select.
        token_counts = tf.cast(positions.row_lengths(), "float32")
        num_to_select = tf.random.stateless_binomial(
            shape=tf.shape(token_counts),
            seed=self._generator.make_seeds()[:, 0],
            counts=token_counts,
            probs=self.rate,
        )
        if self.max_swaps is not None:
            num_to_select = tf.math.minimum(num_to_select, self.max_swaps)
        num_to_select = tf.math.minimum(
            num_to_select, tf.cast(positions.row_lengths(), "int32")
        )
        num_to_select = tf.cast(num_to_select, "int64")

        def _swap(x):
            positions, inputs, num_to_select = x
            for _ in range(num_to_select):
                index = tf.random.stateless_uniform(
                    shape=[2],
                    minval=0,
                    maxval=tf.size(positions),
                    dtype="int32",
                    seed=self._generator.make_seeds()[:, 0],
                )
                index1, index2 = positions[index[0]], positions[index[1]]
                # swap items at the sampled indices with each other
                inputs = tf.tensor_scatter_nd_update(
                    inputs,
                    [[index1], [index2]],
                    [inputs[index2], inputs[index1]],
                )
            return inputs

        swapped = tf.map_fn(
            _swap,
            (positions, inputs, num_to_select),
            fn_output_signature=tf.RaggedTensorSpec(
                ragged_rank=positions.ragged_rank - 1, dtype=inputs.dtype
            ),
        )
        swapped.flat_values.set_shape([None])

        if unbatched:
            swapped = tf.squeeze(swapped, axis=0)
        return swapped
//  [python]
def dice_hard_coe(output, target, threshold=0.5, axis=[1,2,3], smooth=1e-5):
    """Non-differentiable Sørensen–Dice coefficient for comparing the similarity
    of two batch of data, usually be used for binary image segmentation i.e. labels are binary.
    The coefficient between 0 to 1, 1 if totally match.

    Parameters
    -----------
    output : tensor
        A distribution with shape: [batch_size, ....], (any dimensions).
    target : tensor
        A distribution with shape: [batch_size, ....], (any dimensions).
    threshold : float
        The threshold value to be true.
    axis : list of integer
        All dimensions are reduced, default ``[1,2,3]``.
    smooth : float
        This small value will be added to the numerator and denominator, see ``dice_coe``.

    References
    -----------
    - `Wiki-Dice <https://en.wikipedia.org/wiki/Sørensen–Dice_coefficient>`_
    """
    output = tf.cast(output > threshold, dtype=tf.float32)
    target = tf.cast(target > threshold, dtype=tf.float32)
    inse = tf.reduce_sum(tf.multiply(output, target), axis=axis)
    l = tf.reduce_sum(output, axis=axis)
    r = tf.reduce_sum(target, axis=axis)
    ## old axis=[0,1,2,3]
    # hard_dice = 2 * (inse) / (l + r)
    # epsilon = 1e-5
    # hard_dice = tf.clip_by_value(hard_dice, 0, 1.0-epsilon)
    ## new haodong
    hard_dice = (2. * inse + smooth) / (l + r + smooth)
    ##
    hard_dice = tf.reduce_mean(hard_dice)
    return hard_dice
//  [python]
def iou_coe(output, target, threshold=0.5, axis=[1,2,3], smooth=1e-5):
    """Non-differentiable Intersection over Union (IoU) for comparing the
    similarity of two batch of data, usually be used for evaluating binary image segmentation.
    The coefficient between 0 to 1, 1 means totally match.

    Parameters
    -----------
    output : tensor
        A distribution with shape: [batch_size, ....], (any dimensions).
    target : tensor
        A distribution with shape: [batch_size, ....], (any dimensions).
    threshold : float
        The threshold value to be true.
    axis : list of integer
        All dimensions are reduced, default ``[1,2,3]``.
    smooth : float
        This small value will be added to the numerator and denominator, see ``dice_coe``.

    Notes
    ------
    - IoU cannot be used as training loss, people usually use dice coefficient for training, IoU and hard-dice for evaluating.
    """
    pre = tf.to_float(tf.greater(output, threshold))
    truth = tf.to_float(tf.greater(target, threshold))
    inse = tf.reduce_sum(tf.multiply(pre, truth), axis=axis) # AND
    union = tf.reduce_sum(tf.to_float(tf.greater_equal(tf.add(pre, truth), 1)), axis=axis) # OR
    ## old axis=[0,1,2,3]
    # epsilon = 1e-5
    # batch_iou = inse / (union + epsilon)
    ## new haodong
    batch_iou = (inse + smooth) / (union + smooth)
    iou = tf.reduce_mean(batch_iou)
    return iou