# Lint as: python2, python3
# Copyright 2020 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Common utility for object detection tf.train.SequenceExamples."""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import numpy as np
import tensorflow.compat.v1 as tf


def context_float_feature(ndarray):
  """Converts a numpy float array to a context float feature.

  Args:
    ndarray: A numpy float array.

  Returns:
    A context float feature.
  """
  feature = tf.train.Feature()
  for val in ndarray:
    feature.float_list.value.append(val)
  return feature


def context_int64_feature(ndarray):
  """Converts a numpy array to a context int64 feature.

  Args:
    ndarray: A numpy int64 array.

  Returns:
    A context int64 feature.
  """
  feature = tf.train.Feature()
  for val in ndarray:
    feature.int64_list.value.append(val)
  return feature


def context_bytes_feature(ndarray):
  """Converts a numpy bytes array to a context bytes feature.

  Args:
    ndarray: A numpy bytes array.

  Returns:
    A context bytes feature.
  """
  feature = tf.train.Feature()
  for val in ndarray:
    if isinstance(val, np.ndarray):
      val = val.tolist()
    feature.bytes_list.value.append(tf.compat.as_bytes(val))
  return feature


def sequence_float_feature(ndarray):
  """Converts a numpy float array to a sequence float feature.

  Args:
    ndarray: A numpy float array.

  Returns:
    A sequence float feature.
  """
  feature_list = tf.train.FeatureList()
  for row in ndarray:
    feature = feature_list.feature.add()
    if row.size:
      feature.float_list.value[:] = row
  return feature_list


def sequence_int64_feature(ndarray):
  """Converts a numpy int64 array to a sequence int64 feature.

  Args:
    ndarray: A numpy int64 array.

  Returns:
    A sequence int64 feature.
  """
  feature_list = tf.train.FeatureList()
  for row in ndarray:
    feature = feature_list.feature.add()
    if row.size:
      feature.int64_list.value[:] = row
  return feature_list


def sequence_bytes_feature(ndarray):
  """Converts a bytes float array to a sequence bytes feature.

  Args:
    ndarray: A numpy bytes array.

  Returns:
    A sequence bytes feature.
  """
  feature_list = tf.train.FeatureList()
  for row in ndarray:
    if isinstance(row, np.ndarray):
      row = row.tolist()
    feature = feature_list.feature.add()
    if row:
      row = [tf.compat.as_bytes(val) for val in row]
      feature.bytes_list.value[:] = row
  return feature_list


def boxes_to_box_components(bboxes):
  """Converts a list of numpy arrays (boxes) to box components.

  Args:
    bboxes: A numpy array of bounding boxes.

  Returns:
    Bounding box component lists.
  """
  ymin_list = []
  xmin_list = []
  ymax_list = []
  xmax_list = []
  for bbox in bboxes:
    bbox = np.array(bbox).astype(np.float32)
    ymin, xmin, ymax, xmax = np.split(bbox, 4, axis=1)
    ymin_list.append(np.reshape(ymin, [-1]))
    xmin_list.append(np.reshape(xmin, [-1]))
    ymax_list.append(np.reshape(ymax, [-1]))
    xmax_list.append(np.reshape(xmax, [-1]))
  return ymin_list, xmin_list, ymax_list, xmax_list


def make_sequence_example(dataset_name,
                          video_id,
                          encoded_images,
                          image_height,
                          image_width,
                          image_format=None,
                          image_source_ids=None,
                          timestamps=None,
                          is_annotated=None,
                          bboxes=None,
                          label_strings=None,
                          detection_bboxes=None,
                          detection_classes=None,
                          detection_scores=None):
  """Constructs tf.SequenceExamples.

  Args:
    dataset_name: String with dataset name.
    video_id: String with video id.
    encoded_images: A [num_frames] list (or numpy array) of encoded image
      frames.
    image_height: Height of the images.
    image_width: Width of the images.
    image_format: Format of encoded images.
    image_source_ids: (Optional) A [num_frames] list of unique string ids for
      each image.
    timestamps: (Optional) A [num_frames] list (or numpy array) array with image
      timestamps.
    is_annotated: (Optional) A [num_frames] list (or numpy array) array
      in which each element indicates whether the frame has been annotated
      (1) or not (0).
    bboxes: (Optional) A list (with num_frames elements) of [num_boxes_i, 4]
      numpy float32 arrays holding boxes for each frame.
    label_strings: (Optional) A list (with num_frames_elements) of [num_boxes_i]
      numpy string arrays holding object string labels for each frame.
    detection_bboxes: (Optional) A list (with num_frames elements) of
      [num_boxes_i, 4] numpy float32 arrays holding prediction boxes for each
      frame.
    detection_classes: (Optional) A list (with num_frames_elements) of
      [num_boxes_i] numpy int64 arrays holding predicted classes for each frame.
    detection_scores: (Optional) A list (with num_frames_elements) of
      [num_boxes_i] numpy float32 arrays holding predicted object scores for
      each frame.

  Returns:
    A tf.train.SequenceExample.
  """
  num_frames = len(encoded_images)
  image_encoded = np.expand_dims(encoded_images, axis=-1)
  if timestamps is None:
    timestamps = np.arange(num_frames)
  image_timestamps = np.expand_dims(timestamps, axis=-1)

  # Context fields.
  context_dict = {
      'example/dataset_name': context_bytes_feature([dataset_name]),
      'clip/start/timestamp': context_int64_feature([image_timestamps[0][0]]),
      'clip/end/timestamp': context_int64_feature([image_timestamps[-1][0]]),
      'clip/frames': context_int64_feature([num_frames]),
      'image/channels': context_int64_feature([3]),
      'image/height': context_int64_feature([image_height]),
      'image/width': context_int64_feature([image_width]),
      'clip/media_id': context_bytes_feature([video_id])
  }

  # Sequence fields.
  feature_list = {
      'image/encoded': sequence_bytes_feature(image_encoded),
      'image/timestamp': sequence_int64_feature(image_timestamps),
  }

  # Add optional fields.
  if image_format is not None:
    context_dict['image/format'] = context_bytes_feature([image_format])
  if image_source_ids is not None:
    feature_list['image/source_id'] = sequence_bytes_feature(image_source_ids)
  if bboxes is not None:
    bbox_ymin, bbox_xmin, bbox_ymax, bbox_xmax = boxes_to_box_components(bboxes)
    feature_list['region/bbox/xmin'] = sequence_float_feature(bbox_xmin)
    feature_list['region/bbox/xmax'] = sequence_float_feature(bbox_xmax)
    feature_list['region/bbox/ymin'] = sequence_float_feature(bbox_ymin)
    feature_list['region/bbox/ymax'] = sequence_float_feature(bbox_ymax)
    if is_annotated is None:
      is_annotated = np.ones(num_frames, dtype=np.int64)
    is_annotated = np.expand_dims(is_annotated, axis=-1)
    feature_list['region/is_annotated'] = sequence_int64_feature(is_annotated)

  if label_strings is not None:
    feature_list['region/label/string'] = sequence_bytes_feature(
        label_strings)

  if detection_bboxes is not None:
    det_bbox_ymin, det_bbox_xmin, det_bbox_ymax, det_bbox_xmax = (
        boxes_to_box_components(detection_bboxes))
    feature_list['predicted/region/bbox/xmin'] = sequence_float_feature(
        det_bbox_xmin)
    feature_list['predicted/region/bbox/xmax'] = sequence_float_feature(
        det_bbox_xmax)
    feature_list['predicted/region/bbox/ymin'] = sequence_float_feature(
        det_bbox_ymin)
    feature_list['predicted/region/bbox/ymax'] = sequence_float_feature(
        det_bbox_ymax)
  if detection_classes is not None:
    feature_list['predicted/region/label/index'] = sequence_int64_feature(
        detection_classes)
  if detection_scores is not None:
    feature_list['predicted/region/label/confidence'] = sequence_float_feature(
        detection_scores)

  context = tf.train.Features(feature=context_dict)
  feature_lists = tf.train.FeatureLists(feature_list=feature_list)

  sequence_example = tf.train.SequenceExample(
      context=context,
      feature_lists=feature_lists)
  return sequence_example