# Copyright 2018 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.
# ==============================================================================

"""LSTM Meta-architecture definition.

General tensorflow implementation of convolutional Multibox/SSD detection
models with LSTM states, for use on video data.

See https://arxiv.org/abs/1711.06368 for details.
"""
import re
import tensorflow as tf

from object_detection.core import box_list_ops
from object_detection.core import standard_fields as fields
from object_detection.meta_architectures import ssd_meta_arch
from object_detection.utils import ops
from object_detection.utils import shape_utils

slim = tf.contrib.slim


class LSTMMetaArch(ssd_meta_arch.SSDMetaArch):
  """LSTM Meta-architecture definition."""

  def __init__(self,
               is_training,
               anchor_generator,
               box_predictor,
               box_coder,
               feature_extractor,
               encode_background_as_zeros,
               image_resizer_fn,
               non_max_suppression_fn,
               score_conversion_fn,
               classification_loss,
               localization_loss,
               classification_loss_weight,
               localization_loss_weight,
               normalize_loss_by_num_matches,
               hard_example_miner,
               unroll_length,
               target_assigner_instance,
               add_summaries=True):
    super(LSTMMetaArch, self).__init__(
        is_training=is_training,
        anchor_generator=anchor_generator,
        box_predictor=box_predictor,
        box_coder=box_coder,
        feature_extractor=feature_extractor,
        encode_background_as_zeros=encode_background_as_zeros,
        image_resizer_fn=image_resizer_fn,
        non_max_suppression_fn=non_max_suppression_fn,
        score_conversion_fn=score_conversion_fn,
        classification_loss=classification_loss,
        localization_loss=localization_loss,
        classification_loss_weight=classification_loss_weight,
        localization_loss_weight=localization_loss_weight,
        normalize_loss_by_num_matches=normalize_loss_by_num_matches,
        hard_example_miner=hard_example_miner,
        target_assigner_instance=target_assigner_instance,
        add_summaries=add_summaries)
    self._unroll_length = unroll_length

  @property
  def unroll_length(self):
    return self._unroll_length

  @unroll_length.setter
  def unroll_length(self, unroll_length):
    self._unroll_length = unroll_length

  def predict(self, preprocessed_inputs, true_image_shapes, states=None,
              state_name='lstm_state', feature_scope=None):