text:
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from __future__ import unicode_literals

import numpy as np
import tensorflow as tf
import unittest

from onnx_tf.backend import onnx_graph_to_tensorflow_rep
from onnx_tf.common.legacy import legacy_opset_pre_ver
from onnx_tf.common.pooling_helper import py_pool
from onnx import defs
from onnx import helper
from onnx import TensorProto

class TestDynamicShape(unittest.TestCase):
  """ Tests for dynamic shape support
  """

  def _get_rnd_float32(self, low=-1.0, high=1.0, shape=None):
    output = np.random.uniform(low, high, shape)
    if shape == None:
      return np.float32(output)
    else:
      return output.astype(np.float32)

  def _get_rnd_int(self, low, high=None, shape=None, dtype=np.int32):
    return np.random.randint(low, high, size=shape, dtype=dtype)

  def test_compress(self):
    if legacy_opset_pre_ver(9):
      raise unittest.SkipTest(
          "ONNX version {} doesn't support Compress.".format(
              defs.onnx_opset_version()))
    axis = 1
    node_def = helper.make_node("Compress",
                                inputs=['X', 'condition'],
                                outputs=['Y'],
                                axis=axis)
    graph_def = helper.make_graph(
        [node_def],
        name="test_unknown_shape",
        inputs=[
            helper.make_tensor_value_info("X", TensorProto.FLOAT,
                                          [None, None, None]),
            helper.make_tensor_value_info("condition", TensorProto.BOOL, [None])
        ],
        outputs=[
            helper.make_tensor_value_info("Y", TensorProto.FLOAT,
                                          [None, None, None])
        ])
    x = self._get_rnd_float32(shape=[5, 5, 5])
    cond = np.array([1, 0, 1])
    tf_rep = onnx_graph_to_tensorflow_rep(graph_def)
    output = tf_rep.run({"X": x, "condition": cond})
    np.testing.assert_almost_equal(output['Y'], np.compress(cond, x, axis=axis))

  def test_eye_like(self):
    if legacy_opset_pre_ver(9):
      raise unittest.SkipTest("ONNX version {} doesn't support EyeLike.".format(
          defs.onnx_opset_version()))
    shape = [6, 10]
    off_diagonal_offset = -3
    x = self._get_rnd_int(0, 100, shape=shape)
    y = np.eye(shape[0], shape[1], k=off_diagonal_offset, dtype=np.float32)
    node_def = helper.make_node("EyeLike", ["x"], ["y"],
                                dtype=TensorProto.FLOAT,
                                k=off_diagonal_offset)
    graph_def = helper.make_graph(
        [node_def],
        name="test_unknown_shape",
        inputs=[
            helper.make_tensor_value_info("x", TensorProto.INT32, [None, None])
        ],
        outputs=[
            helper.make_tensor_value_info("y", TensorProto.FLOAT, [None, None])
        ])
    tf_rep = onnx_graph_to_tensorflow_rep(graph_def)
    output = tf_rep.run({"x": x})
    np.testing.assert_equal(output["y"], y)

  def test_flatten(self):
    shape = [2, 3, 4]
    x = self._get_rnd_float32(shape=shape)
    axis = 1
    node_def = helper.make_node("Flatten", ["X"], ["Y"], axis=axis)
    graph_def = helper.make_graph(
        [node_def],
        name="test_unknown_shape",
        inputs=[
            helper.make_tensor_value_info("X", TensorProto.FLOAT,
                                          [None, None, None])
        ],
        outputs=[helper.make_tensor_value_info("Y", TensorProto.FLOAT, [None])])
    tf_rep = onnx_graph_to_tensorflow_rep(graph_def)
    output = tf_rep.run({"X": x})
    new_shape = (np.prod(shape[0:axis]).astype(int), -1)
    np.testing.assert_almost_equal(output["Y"], np.reshape(x, new_shape))

  def test_gather_nd(self):
    if legacy_opset_pre_ver(11):
      raise unittest.SkipTest(
          "ONNX version {} doesn't support GatherND.".format(
              defs.onnx_opset_version()))
    # valid positive and negative indices for elements
    data = np.array([[1, 2, 3], [4, 5, 6]], dtype=np.int32)
    indices = np.array([[0, 0], [1, -3]], dtype=np.int64)
    ref_output = np.array([1, 4], dtype=np.int32)
    node_def = helper.make_node("GatherND", ["data", "indices"], ["outputs"])
    graph_def = helper.make_graph(
        [node_def],
        name="test_unknown_shape",
        inputs=[
            helper.make_tensor_value_info("data", TensorProto.INT32,
                                          [None, None]),
            helper.make_tensor_value_info("indices", TensorProto.INT64,
                                          [None, None])
        ],
        outputs=[
            helper.make_tensor_value_info("outputs", TensorProto.INT32, [None])
        ])
    tf_rep = onnx_graph_to_tensorflow_rep(graph_def)
    output = tf_rep.run({"data": data, "indices": indices})
    np.testing.assert_almost_equal(output["outputs"], ref_output)

  def test_is_inf(self):
    if legacy_opset_pre_ver(10):
      raise unittest.SkipTest("ONNX version {} doesn't support IsInf.".format(
          defs.onnx_opset_version()))
    inp = np.array([-1.2, np.nan, np.inf, 2.8, np.NINF, np.inf],
                   dtype=np.float32)
    expected_output = np.isinf(inp)
    node_def = helper.make_node("IsInf", ["X"], ["Y"])
    graph_def = helper.make_graph(
        [node_def],
        name="test_unknown_shape",
        inputs=[
            helper.make_tensor_value_info("X", TensorProto.FLOAT, [None]),
        ],
        outputs=[helper.make_tensor_value_info("Y", TensorProto.BOOL, [None])])
    tf_rep = onnx_graph_to_tensorflow_rep(graph_def)
    output = tf_rep.run({"X": inp})
    np.testing.assert_equal(output["Y"], expected_output)

  def test_matmul_integer(self):
    if legacy_opset_pre_ver(10):
      raise unittest.SkipTest(
          "ONNX version {} doesn't support MatMulInteger.".format(
              defs.onnx_opset_version()))

    node_def = helper.make_node("MatMulInteger",
                                ["A", "B", "a_zero_point", "b_zero_point"],
                                ["Z"])
    # A & B are 3-D tensor and a_zero_point & b_zero_point are scalar
    A = self._get_rnd_int(-20, 20, shape=(2, 3, 4), dtype=np.int8)
    B = self._get_rnd_int(-20, 20, shape=(2, 4, 6), dtype=np.int8)
    a_zero_point = self._get_rnd_int(-20, 20, dtype=np.int8)
    b_zero_point = self._get_rnd_int(-20, 20, dtype=np.int8)
    A_minus_zero_point = np.subtract(A.astype(np.int32),
                                     a_zero_point.astype(np.int32))
    B_minus_zero_point = np.subtract(B.astype(np.int32),
                                     b_zero_point.astype(np.int32))
    z = np.matmul(A_minus_zero_point, B_minus_zero_point)
    graph_def = helper.make_graph(
        [node_def],
        name="test_unknown_shape",
        inputs=[
            helper.make_tensor_value_info("A", TensorProto.INT8,
                                          [None, None, None]),
            helper.make_tensor_value_info("B", TensorProto.INT8,
                                          [None, None, None]),
            helper.make_tensor_value_info("a_zero_point", TensorProto.INT8, []),
            helper.make_tensor_value_info("b_zero_point", TensorProto.INT8, [])
        ],
        outputs=[
            helper.make_tensor_value_info("Z", TensorProto.INT32,
                                          [None, None, None])
        ])
    tf_rep = onnx_graph_to_tensorflow_rep(graph_def)
    output = tf_rep.run({
        "A": A,
        "B": B,
        "a_zero_point": a_zero_point,
        "b_zero_point": b_zero_point
    })
    np.testing.assert_almost_equal(output["Z"], z)
    # A & B are 4-D tensor and a_zero_point & b_zero_point are 1-D tensor
    A = self._get_rnd_int(-20, 20, shape=(2, 5, 3, 4), dtype=np.int8)
    B = self._get_rnd_int(-20, 20, shape=(2, 1, 4, 6), dtype=np.int8)
    a_zero_point = self._get_rnd_int(-20,
                                     20,
                                     shape=(A.shape[-2]),
                                     dtype=np.int8)
    b_zero_point = self._get_rnd_int(-20,
                                     20,
                                     shape=(B.shape[-1]),
                                     dtype=np.int8)
    a_zero_point_with_reshape = np.reshape(a_zero_point, [A.shape[-2], 1])
    A_minus_zero_point = np.subtract(A.astype(np.int32),
                                     a_zero_point_with_reshape.astype(np.int32))
    B_minus_zero_point = np.subtract(B.astype(np.int32),
                                     b_zero_point.astype(np.int32))
    z = np.matmul(A_minus_zero_point, B_minus_zero_point)
    graph_def = helper.make_graph(
        [node_def],
        name="test_unknown_shape",
        inputs=[
            helper.make_tensor_value_info("A", TensorProto.INT8,
                                          [None, None, None, None]),
            helper.make_tensor_value_info("B", TensorProto.INT8,
                                          [None, None, None, None]),
            helper.make_tensor_value_info("a_zero_point", TensorProto.INT8,
                                          [None]),
            helper.make_tensor_value_info("b_zero_point", TensorProto.INT8,
                                          [None])
        ],
        outputs=[
            helper.make_tensor_value_info("Z", TensorProto.INT32,
                                          [None, None, None, None])
        ])
    tf_rep = onnx_graph_to_tensorflow_rep(graph_def)
    output = tf_rep.run({
        "A": A,
        "B": B,
        "a_zero_point": a_zero_point,
        "b_zero_point": b_zero_point
    })
    np.testing.assert_almost_equal(output["Z"], z)

  def test_non_max_suppression(self):
    if legacy_opset_pre_ver(10):
      raise unittest.SkipTest(
          "ONNX version {} doesn't support NonMaxSuppression.".format(
              defs.onnx_opset_version()))
    boxes = np.array([[[0.0, 0.0, 1.0, 1.0], [0.0, 0.1, 1.0, 1.1],
                       [0.0, -0.1, 1.0, 0.9], [0.0, 10.0, 1.0, 11.0],
                       [0.0, 10.1, 1.0, 11.1], [0.0, 100.0, 1.0, 101.0]],
                      [[0.0, 0.0, 1.0, 1.0], [0.0, 0.1, 1.0, 1.1],
                       [0.0, -0.1, 1.0, 0.9], [0.0, 10.0, 1.0, 11.0],
                       [0.0, 10.1, 1.0, 11.1], [0.0, 100.0, 1.0,
                                                101.0]]]).astype(np.float32)
    scores = np.array([[[0.9, 0.75, 0.6, 0.95, 0.5, 0.3]],
                       [[0.9, 0.75, 0.6, 0.95, 0.5, 0.3]]]).astype(np.float32)
    max_output_boxes_per_class = np.array([2]).astype(np.int64)
    iou_threshold = np.array([0.5]).astype(np.float32)
    score_threshold = np.array([0.0]).astype(np.float32)
    selected_indices = np.array([[0, 0, 3], [0, 0, 0], [1, 0, 3],
                                 [1, 0, 0]]).astype(np.int64)
    node_def = helper.make_node("NonMaxSuppression", [
        "boxes", "scores", "max_output_boxes_per_class", "iou_threshold",
        "score_threshold"
    ], ["selected_indices"],
                                center_point_box=0)
    graph_def = helper.make_graph(
        [node_def],
        name="test_unknown_shape",
        inputs=[
            helper.make_tensor_value_info("boxes", TensorProto.FLOAT,
                                          [None, None, None]),
            helper.make_tensor_value_info("scores", TensorProto.FLOAT,
                                          [None, None, None]),
            helper.make_tensor_value_info("max_output_boxes_per_class",
                                          TensorProto.INT64, [None]),
            helper.make_tensor_value_info("iou_threshold", TensorProto.FLOAT,
                                          [None]),
            helper.make_tensor_value_info("score_threshold", TensorProto.FLOAT,
                                          [None])
        ],
        outputs=[
            helper.make_tensor_value_info("selected_indices", TensorProto.INT64,
                                          [None, None])
        ])
    tf_rep = onnx_graph_to_tensorflow_rep(graph_def)
    output = tf_rep.run({
        "boxes": boxes,
        "scores": scores,
        "max_output_boxes_per_class": max_output_boxes_per_class,
        "iou_threshold": iou_threshold,
        "score_threshold": score_threshold
    })
    np.testing.assert_almost_equal(output["selected_indices"], selected_indices)

  def test_scatter_nd(self):
    if legacy_opset_pre_ver(11):
      raise unittest.SkipTest(
          "ONNX version {} doesn't support ScatterND.".format(
              defs.onnx_opset_version()))
    # valid positive and negative indices for slices
    data = np.reshape(np.arange(1, 25, dtype=np.float32), [2, 3, 4])
    indices = np.array([[-1]], dtype=np.int64)
    updates = np.array([[[43, 44, 45, 46], [47, 48, 49, 50], [51, 52, 53, 54]]],
                       dtype=np.float32)
    ref_output = np.array(
        [[[1, 2, 3, 4], [5, 6, 7, 8], [9, 10, 11, 12]],
         [[43, 44, 45, 46], [47, 48, 49, 50], [51, 52, 53, 54]]],
        dtype=np.float32)
    node_def = helper.make_node("ScatterND", ["data", "indices", "updates"],
                                ["outputs"])
    graph_def = helper.make_graph(
        [node_def],
        name="test_unknown_shape",
        inputs=[
            helper.make_tensor_value_info("data", TensorProto.FLOAT,
                                          [None, None, None]),
            helper.make_tensor_value_info("indices", TensorProto.INT64,
                                          [None, None]),
            helper.make_tensor_value_info("updates", TensorProto.FLOAT,
                                          [None, None, None])
        ],
        outputs=[
            helper.make_tensor_value_info("outputs", TensorProto.FLOAT,
                                          [None, None, None])
        ])
    tf_rep = onnx_graph_to_tensorflow_rep(graph_def)
    output = tf_rep.run({"data": data, "indices": indices, "updates": updates})
    np.testing.assert_almost_equal(output["outputs"], ref_output)

  def test_max_pool_2d_dilations_ceil_pads(self):
    if legacy_opset_pre_ver(10):
      raise unittest.SkipTest(
          "ONNX version {} doesn't support dilations nor ceil mode.".format(
              defs.onnx_opset_version()))

    kernel_shape = [3, 3]
    strides = [2, 2]
    dilations = [3, 3]
    pads = [1, 1, 2, 2]
    ceil_mode = 1

    input_shape = [10, 3, 23, 23]
    x = self._get_rnd_float32(shape=input_shape)

    test_output = py_pool(x,
                          kernel_shape=kernel_shape,
                          strides=strides,
                          dilations=dilations,
                          padding=pads,
                          ceil_mode=ceil_mode,
                          pooling_type="MAX",
                          include_indices=False)

    node_def = helper.make_node(op_type="MaxPool",
                                inputs=["X"],
                                outputs=["Y"],
                                kernel_shape=kernel_shape,
                                strides=strides,
                                dilations=dilations,
                                pads=pads,
                                ceil_mode=ceil_mode)

    graph_def = helper.make_graph(
        [node_def],
        name="test_unknown_shape",
        inputs=[
            helper.make_tensor_value_info("X", TensorProto.FLOAT,
                                          [None, None, None, None]),
        ],
        outputs=[
            helper.make_tensor_value_info("Y", TensorProto.FLOAT,
                                          [None, None, None, None])
        ])
    tf_rep = onnx_graph_to_tensorflow_rep(graph_def)
    output = tf_rep.run({"X": x})

    np.testing.assert_almost_equal(output["Y"], test_output)

  def test_average_pool_2d(self):
    kernel_shape = [1, 2]
    strides = [1, 2]

    input_shape = [10, 10, 4, 4]
    x = self._get_rnd_float32(shape=input_shape)

    test_output = py_pool(x,
                          kernel_shape=kernel_shape,
                          strides=strides,
                          pooling_type="AVG",
                          include_indices=False)

    node_def = helper.make_node(op_type="AveragePool",
                                inputs=["X"],
                                outputs=["Y"],
                                kernel_shape=kernel_shape,
                                strides=strides)

    graph_def = helper.make_graph(
        [node_def],
        name="test_unknown_shape",
        inputs=[
            helper.make_tensor_value_info("X", TensorProto.FLOAT,
                                          [None, None, None, None]),
        ],
        outputs=[
            helper.make_tensor_value_info("Y", TensorProto.FLOAT,
                                          [None, None, None, None])
        ])
    tf_rep = onnx_graph_to_tensorflow_rep(graph_def)
    output = tf_rep.run({"X": x})

    np.testing.assert_almost_equal(output["Y"], test_output)

  def test_max_unpool(self):
    input_shape = [10, 3, 24, 24]
    x = self._get_rnd_float32(shape=input_shape)

    kernel_shape = [2, 2]
    strides = [2, 2]

    maxpool_node_def = helper.make_node(
            op_type="MaxPool",
            inputs=["X"],
            outputs=["Pool", "Indices"],
            kernel_shape=kernel_shape,
            strides=strides)

    maxunpool_node_def = helper.make_node(
        "MaxUnpool", ["Pool", "Indices"], ["Y"],
        kernel_shape=kernel_shape,
        strides=strides)

    graph_def = helper.make_graph(
        [maxpool_node_def,maxunpool_node_def],
        name="test_unknown_shape",
        inputs=[
            helper.make_tensor_value_info("X", TensorProto.FLOAT,
                                          [None, None, None, None]),
        ],
        outputs=[
            helper.make_tensor_value_info("Y", TensorProto.FLOAT,
                                          [None, None, None, None])
        ])
    tf_rep = onnx_graph_to_tensorflow_rep(graph_def)
    output_unpool = tf_rep.run({"X": x})

    test_output = np.zeros(input_shape)
    for i1 in range(0, input_shape[0]):
      for i2 in range(0, input_shape[1]):
        for i3 in range(0, input_shape[2], 2):
          for i4 in range(0, input_shape[3], 2):
            max_val = float('-inf')
            for j1 in range(i3,i3+2):
              for j2 in range(i4,i4+2):
                if x[i1][i2][j1][j2] > max_val:
                  max_val = x[i1][i2][j1][j2]
                  max_ind = (j1, j2)
            j1, j2 = max_ind
            test_output[i1][i2][j1][j2] = max_val
    np.testing.assert_almost_equal(output_unpool["Y"], test_output)

  def test_slice(self):
    # test case 1 with normal inputs
    axes = [0, 1, 2]
    starts = [0, 0, 0]
    ends = [2, 2, 2]

    if legacy_opset_pre_ver(10):
      node_def = helper.make_node("Slice", ["X"], ["S"],
                                  axes=axes,
                                  starts=starts,
                                  ends=ends)
      graph_def = helper.make_graph(
          [node_def],
          name="test_unknown_shape",
          inputs=[
              helper.make_tensor_value_info("X", TensorProto.FLOAT,
                                            [None, None, None])
          ],
          outputs=[
              helper.make_tensor_value_info("S", TensorProto.FLOAT,
                                            [None, None, None])
          ])
    else:
      node_def = helper.make_node("Slice", ["X", "starts", "ends", "axes"],
                                  ["S"])
      graph_def = helper.make_graph(
          [node_def],
          name="test_unknown_shape",
          inputs=[
              helper.make_tensor_value_info("X", TensorProto.FLOAT,
                                            [None, None, None]),
              helper.make_tensor_value_info("starts", TensorProto.INT32,
                                            [None]),
              helper.make_tensor_value_info("ends", TensorProto.INT32, [None]),
              helper.make_tensor_value_info("axes", TensorProto.INT32, [None]),
          ],
          outputs=[
              helper.make_tensor_value_info("S", TensorProto.FLOAT,
                                            [None, None, None])
          ])
    tf_rep = onnx_graph_to_tensorflow_rep(graph_def)

    if legacy_opset_pre_ver(10):
      x = self._get_rnd_float32(shape=[1000]).reshape([10, 10, 10])
      output = tf_rep.run({"X": x})
      np.testing.assert_almost_equal(output["S"], x[0:2, 0:2, 0:2])
    else:
      x = self._get_rnd_float32(shape=[1000]).reshape([10, 10, 10])
      output = tf_rep.run({
          "X": x,
          "starts": starts,
          "ends": ends,
          "axes": axes
      })
      np.testing.assert_almost_equal(output["S"], x[0:2, 0:2, 0:2])

    # test case 2 with negative, out-of-bound and default inputs
    axes = [0, 2]
    starts = [0, -7]
    ends = [-8, 20]
    steps = [1, 1]

    if legacy_opset_pre_ver(10):
      node_def = helper.make_node("Slice", ["X"], ["S"],
                                  axes=axes,
                                  starts=starts,
                                  ends=ends)
      graph_def = helper.make_graph(
          [node_def],
          name="test_unknown_shape",
          inputs=[
              helper.make_tensor_value_info("X", TensorProto.FLOAT,
                                            [None, None, None])
          ],
          outputs=[
              helper.make_tensor_value_info("S", TensorProto.FLOAT,
                                            [None, None, None])
          ])
    else:
      node_def = helper.make_node("Slice",
                                  ["X", "starts", "ends", "axes", "steps"],
                                  ["S"])
      graph_def = helper.make_graph(
          [node_def],
          name="test_unknown_shape",
          inputs=[
              helper.make_tensor_value_info("X", TensorProto.FLOAT,
                                            [None, None, None]),
              helper.make_tensor_value_info("starts", TensorProto.INT32,
                                            [None]),
              helper.make_tensor_value_info("ends", TensorProto.INT32, [None]),
              helper.make_tensor_value_info("axes", TensorProto.INT32, [None]),
              helper.make_tensor_value_info("steps", TensorProto.INT32, [None]),
          ],
          outputs=[
              helper.make_tensor_value_info("S", TensorProto.FLOAT,
                                            [None, None, None])
          ])
    tf_rep = onnx_graph_to_tensorflow_rep(graph_def)
    if legacy_opset_pre_ver(10):
      x = self._get_rnd_float32(shape=[1000]).reshape([10, 10, 10])
      output = tf_rep.run({"X": x})
      np.testing.assert_almost_equal(output["S"], x[0:-8, :, -7:20])
    else:
      x = self._get_rnd_float32(shape=[1000]).reshape([10, 10, 10])
      output = tf_rep.run({
          "X": x,
          "starts": starts,
          "ends": ends,
          "axes": axes,
          "steps": steps
      })
      np.testing.assert_almost_equal(output["S"], x[0:-8, :, -7:20])

    # test case 3 with non-default steps
    axes = [0, 1, 2]
    starts = [0, 0, 0]
    ends = [2, 2, 2]
    steps = [2, -2, -1]

    if not legacy_opset_pre_ver(10):
      x = self._get_rnd_float32(shape=[1000]).reshape([10, 10, 10])
      output = tf_rep.run({
          "X": x,
          "starts": starts,
          "ends": ends,
          "axes": axes,
          "steps": steps
      })
      np.testing.assert_almost_equal(output["S"], x[0:2:2, 0:2:-2, 0:2:-1])

  def test_split(self):
    shape = [12, 12]
    axis = 0
    output_count = 3
    node_def = helper.make_node("Split", ["X"],
                                ["Z%i" % i for i in range(output_count)],
                                axis=axis)
    graph_def = helper.make_graph(
        [node_def],
        name="test_unknown_shape",
        inputs=[
            helper.make_tensor_value_info("X", TensorProto.FLOAT,
                                          [None] * len(shape))
        ],
        outputs=[
            helper.make_tensor_value_info("Z%i" % i, TensorProto.FLOAT,
                                          [None] * len(shape))
            for i in range(output_count)
        ])

    tf_rep = onnx_graph_to_tensorflow_rep(graph_def)
    x = self._get_rnd_float32(shape=shape)
    output = tf_rep.run({"X": x})

    per_part = shape[axis] // output_count
    split = [per_part] * output_count
    for a, b in zip(list(output), np.split(x, np.cumsum(split))[:-1]):
      np.testing.assert_almost_equal(a, b)

if __name__ == '__main__':
  unittest.main()

max_stars_repo_name: heliumsun/onnx-tensorflow

max_issues_repo_name: heliumsun/onnx-tensorflow

max_forks_repo_name: heliumsun/onnx-tensorflow


text:
# -*- coding: utf-8 -*-

"""Behavior manager for platooning specifically
"""
# Author: Runsheng Xu <rxx3386@ucla.edu>
# License: MIT

import weakref
from collections import deque

import carla
import numpy as np

from opencda.core.application.platooning.fsm import FSM
from opencda.core.application.platooning.platoon_debug_helper import PlatoonDebugHelper
from opencda.core.common.misc import compute_distance, get_speed, cal_distance_angle
from opencda.core.plan.behavior_agent import BehaviorAgent

class PlatooningBehaviorAgent(BehaviorAgent):
    """
    Platoon behavior agent that inherits the single vehicle behavior agent.

    Parameters
    ----------
    vehicle : carla.Vehicle
        The carla vehicle.
    vehicle_manager : opencda object
        The vehicle manager, used when joining platoon finished.
    v2x_manager : opencda object
        Used to received and deliver information.
    behavior_yaml : dict
        The configuration dictionary for BehaviorAgent.
    platoon_yaml : dict.
        The configuration dictionary for platoon behavior.
    carla_map : carla.Map
        The HD Map used in the simulation.

    Attributes
    ----------
    vehicle_manager : opencda object
        The weak reference of the vehicle manager, used when joining platoon finished.
    v2x_manager : opencda object
        The weak reference of the v2x_manager
    debug_helper : opencda Object
        A debug helper used to record the driving performance during platooning.
    inter_gap : float
        The desired time gap between each platoon member.
    """

    def __init__(self, vehicle, vehicle_manager, v2x_manager, behavior_yaml, platoon_yaml, carla_map):

        super(PlatooningBehaviorAgent, self).__init__(vehicle, carla_map, behavior_yaml)

        self.vehicle_manager = weakref.ref(vehicle_manager)()
        # communication manager
        self.v2x_manager = weakref.ref(v2x_manager)()

        # used for gap keeping
        self.inter_gap = platoon_yaml['inter_gap']
        # used when open a gap
        self.open_gap = platoon_yaml['open_gap']
        # this is used to control gap opening during cooperative joining
        self.current_gap = self.inter_gap

        # used for merging vehicle
        self.destination_changed = False

        # merging vehicle needs to reach this speed before cooperative merge
        self.warm_up_speed = platoon_yaml['warm_up_speed']

        # used to calculate performance
        self.debug_helper = PlatoonDebugHelper(self.vehicle.id)
        self.time_gap = 100.0
        self.dist_gap = 100.0

    def run_step(self, target_speed=None, collision_detector_enabled=True, lane_change_allowed=True):
        """
        Run a single step for navigation under platooning agent. Finite state machine is used to switch between
        different platooning states.

        Parameters
        ----------
        target_speed : float
            Target speed in km/h
        collision_detector_enabled : bool
            Whether collision detection enabled.
        lane_change_allowed : bool
            Whether lane change is allowed.
        Returns
        -------
        """
        # reset time gap and distance gap record at the beginning
        self.time_gap = 100.0
        self.dist_gap = 100.0

        status = self.v2x_manager.get_platoon_status()
        # case1: the vehicle is not cda enabled
        if status == FSM.DISABLE:
            return super().run_step(target_speed, collision_detector_enabled)

        # case2: single vehicle keep searching platoon to join
        if status == FSM.SEARCHING:
            find_platoon, min_index, white_list = self.v2x_manager.match_platoon()
            # we only ignore platoon members for collision detection during joining process
            self.white_list = white_list

            # no platoon found, stay in searching status
            if not find_platoon:
                return super().run_step(target_speed, collision_detector_enabled)

            # platoon found and agreement achieved
            front_vehicle, rear_vehicle = self.v2x_manager.get_platoon_front_rear()
            # if no front vehicle, meaning it will be a frontal joining
            if not front_vehicle and rear_vehicle:
                print('merging vehicle chooses frontal joining')
                self.v2x_manager.set_platoon_status(FSM.FRONT_JOINING)
            # if front vehicle and rear vehicle both exist
            if front_vehicle and rear_vehicle:
                print("merging vehicle chooses cut-in joining")
                self.v2x_manager.set_platoon_status(FSM.MOVE_TO_POINT)
            # if only front vehicle exits
            if front_vehicle and not rear_vehicle:
                print("merging vehicle chooses back joining")
                self.v2x_manager.set_platoon_status(FSM.BACK_JOINING)

            return super().run_step(target_speed, collision_detector_enabled)

        # case3.1: the merging vehicle chooses cut-in joining and is moving to the meeting point
        if status == FSM.MOVE_TO_POINT:
            target_speed, target_waypoint, new_status = self.run_step_cut_in_move2point()
            self.v2x_manager.set_platoon_status(new_status)
            return target_speed, target_waypoint
        # case3.2: the merging vehicle chooses cut-in joining and is ready for merging
        if status == FSM.JOINING:
            target_speed, target_waypoint, new_status = self.run_step_cut_in_joining()
            # if joining is finished
            if new_status == FSM.JOINING_FINISHED:
                self.joining_finish_manager()
            return target_speed, target_waypoint

        # case 4: the merging vehicle selects back joining
        if status == FSM.BACK_JOINING:
            target_speed, target_waypoint, new_status = self.run_step_back_joining()
            # if joining is finshed
            if new_status == FSM.JOINING_FINISHED:
                self.joining_finish_manager()
            return target_speed, target_waypoint

        # case 5: the merging vehicle selects frontal joining
        if status == FSM.FRONT_JOINING:
            target_speed, target_waypoint, new_status = self.run_step_front_joining()
            self.v2x_manager.set_platoon_status(new_status)

            # if joining abandoned
            if new_status == FSM.ABONDON:
                self.v2x_manager.set_platoon_status(FSM.SEARCHING)
                _, rear_vehicle_manager = self.v2x_manager.get_platoon_front_rear()
                self.v2x_manager.add_platoon_blacklist(rear_vehicle_manager.v2x_manager.get_platoon_manager()[0].pmid)
            if new_status == FSM.JOINING_FINISHED:
                self.joining_finish_manager('rear')

            return target_speed, target_waypoint

        # case 6: leading vehicle behavior
        if status == FSM.LEADING_MODE:
            return super().run_step(target_speed, collision_detector_enabled)

        # case7: maintaining status
        if status == FSM.MAINTINING:
            return self.run_step_maintaining()

        # case8: Open Gap status
        if status == FSM.OPEN_GAP:
            return self.run_step_open_gap()

    def update_information(self, ego_pos, ego_speed, objects):
        """
        Update the perception and localization information to the behavior agent.

        Parameters
        ----------
        ego_pos : carla.Transform
            ego position from localization module.
        ego_speed : float
            km/h, ego speed.
        objects : dictionary
            Objects detection results from perception module.

        Returns
        -------
        """
        # update localization information
        self._ego_speed = ego_speed
        self._ego_pos = ego_pos
        self.break_distance = self._ego_speed / 3.6 * self.emergency_param
        # update the localization info to trajectory planner
        self.get_local_planner().update_information(ego_pos, ego_speed)

        # current version only consider about vehicles
        obstacle_vehicles = objects['vehicles']
        self.obstacle_vehicles = self.white_list_match(obstacle_vehicles)

        # update the debug helper
        self.debug_helper.update(ego_speed, self.ttc, time_gap=self.time_gap, dist_gap=self.dist_gap)

        if self.ignore_traffic_light:
            self.light_state = "Green"
        else:
            # This method also includes stop signs and intersections.
            self.light_state = str(self.vehicle.get_traffic_light_state())

    def joining_finish_manager(self, insert_vehicle='front'):
        """
        Called when a joining is finish to update the platoon manager list.

        Parameters
        ----------
        insert_vehicle : string
            Indicate use the front or rear vehicle index to update the platoon manager list.
        """
        frontal_vehicle_manager, rear_vehicle_manager = self.v2x_manager.get_platoon_front_rear()
        if insert_vehicle == 'front':
            platoon_manger, index = frontal_vehicle_manager.v2x_manager.get_platoon_manager()
            platoon_manger.set_member(self.vehicle_manager, index + 1)
            self._local_planner.debug_trajectory = False
        else:
            platoon_manger, index = rear_vehicle_manager.v2x_manager.get_platoon_manager()
            platoon_manger.set_member(self.vehicle_manager, index, lead=True)

        platoon_manger.update_member_order()

    def calculate_gap(self, distance):
        """
        Calculate the current vehicle and frontal vehicle's time/distance gap.

        Args:
            -distance (float):  distance between the ego vehicle and frontal vehicle
        """
        # we need to count the vehicle length in to calculate the gap
        boundingbox = self.vehicle.bounding_box
        veh_length = 2 * abs(boundingbox.location.y - boundingbox.extent.y)

        delta_v = self._ego_speed / 3.6
        time_gap = distance / delta_v
        self.time_gap = time_gap
        self.dist_gap = distance - veh_length

    def platooning_following_manager(self, inter_gap):
        """
        Car following behavior in platooning with gap regulation.

        Args:
            - inter_gap (float): the gap designed for platooning
        """

        frontal_vehicle_manager, _ = self.v2x_manager.get_platoon_front_rear()
        frontal_front_vehicle_manger, _ = frontal_vehicle_manager.v2x_manager.get_platoon_front_rear()

        if len(self._local_planner.get_trajetory()) > self.get_local_planner().trajectory_update_freq - 2:
            return self._local_planner.run_step([], [], [], following=True)
        else:
            # this agent is a behavior agent
            frontal_trajectory = frontal_vehicle_manager.agent.get_local_planner().get_trajetory()

            # get front speed
            frontal_speed = frontal_vehicle_manager.agent._ego_speed

            ego_trajetory = deque(maxlen=30)
            ego_loc_x, ego_loc_y, ego_loc_z = self._ego_pos.location.x, self._ego_pos.location.y, \
                                              self._ego_pos.location.z

            # get ego speed
            ego_speed = self._ego_speed

            # compare speed with frontal veh
            frontal_speedd_diff = ego_speed - frontal_speed

            tracked_length = len(frontal_trajectory) - 1 if not frontal_front_vehicle_manger \
                else len(frontal_trajectory)

            # todo: current not working well on curve
            for i in range(tracked_length):
                delta_t = self.get_local_planner().dt
                # if leader is slowing down(leader target speed is smaller than current speed), use a bigger dt.
                # spd diff max at 15. If diff greater than 8, increase dt
                if frontal_speedd_diff > 3.0:
                    '''
                    # only increase dt when V_ego > V_front (avoid collision)
                    # if V_ego < V_front (diff < 0), stick with small dt
                    # todo: change delta_t to a function:
                    #      --> 1. {V_ego > V_front}: decrease dt to increase gap, help avoid collision
                    #      --> 2. more difference, more dt adjustment
                    #      --> 3. {V_ego < V_front}: will not collide, keep default dt to keep gap
                    #      --> 4. {V_ego ~ V_front}: keep default dt to keep gap
                    '''
                    delta_t = delta_t + frontal_speedd_diff * 0.0125

                # print('previous x :%f, delta t: %f' % (frontal_trajectory[i][0].location.x, delta_t))
                if i == 0:
                    pos_x = (frontal_trajectory[i][0].location.x + inter_gap / delta_t * ego_loc_x) / \
                            (1 + inter_gap / delta_t)
                    pos_y = (frontal_trajectory[i][0].location.y + inter_gap / delta_t * ego_loc_y) / \
                            (1 + inter_gap / delta_t)
                else:
                    pos_x = (frontal_trajectory[i][0].location.x +
                             inter_gap / delta_t * ego_trajetory[i - 1][0].location.x) / \
                            (1 + inter_gap / delta_t)
                    pos_y = (frontal_trajectory[i][0].location.y +
                             inter_gap / delta_t * ego_trajetory[i - 1][0].location.y) / \
                            (1 + inter_gap / delta_t)

                distance = np.sqrt((pos_x - ego_loc_x) ** 2 + (pos_y - ego_loc_y) ** 2)
                velocity = distance / delta_t * 3.6
                ego_trajetory.append([carla.Transform(
                    carla.Location(pos_x, pos_y,
                                   self._map.get_waypoint(self._ego_pos.location).transform.location.z)),
                    velocity])

                ego_loc_x = pos_x
                ego_loc_y = pos_y

            if not ego_trajetory:
                wpt = self._map.get_waypoint(self._ego_pos.location)
                next_wpt = wpt.next(max(2, int(self._ego_speed / 3.6 * 1)))[0]
                ego_trajetory.append((next_wpt.transform,
                                      self._ego_speed))

            return self._local_planner.run_step([], [], [], trajectory=ego_trajetory)

    def platooning_merge_management(self, frontal_vehicle_vm):
        """
        Merge the vehicle into the platooning.

        Parameters
        ----------
        frontal_vehicle_vm : opencda object
            The vehivlel manager of the front vehicle.

        Returns
        -------
        target_speed : float
            The target speed for ego vehicle.
        target_waypoint : carla.waypoint
            The target waaypoint for ego vehcile.
        """
        print("start merging !")
        self.lane_change_allowed = True
        frontal_vehicle_loc = frontal_vehicle_vm.vehicle.get_location()

        # we choose next waypoint of the frontal vehicle as starting point to have smooth speed
        frontal_vehicle_waypoint = frontal_vehicle_vm.agent._map.get_waypoint(frontal_vehicle_loc)
        frontal_vehicle_next_waypoint = frontal_vehicle_waypoint.next(
            get_speed(frontal_vehicle_vm.vehicle, True))[0].transform.location

        # retrieve the platooning's destination
        platooning_manager, _ = frontal_vehicle_vm.v2x_manager.get_platoon_manager()
        destination = platooning_manager.destination

        # regenerate route the route to make merge(lane change)
        self.set_destination(frontal_vehicle_next_waypoint, destination, clean=True)

        target_speed, target_waypoint = super().run_step(target_speed=1.5 * get_speed(frontal_vehicle_vm.vehicle),
                                                         collision_detector_enabled=False)

        return target_speed, target_waypoint

    def run_step_maintaining(self):
        """
        Next step behavior planning for speed maintaining.
        """
        frontal_vehicle_manager, _ = self.v2x_manager.get_platoon_front_rear()
        self.current_gap = self.inter_gap

        frontal_vehicle = frontal_vehicle_manager.vehicle
        frontal_vehicle_loc = frontal_vehicle.get_location()
        ego_vehicle_loc = self._ego_pos.location

        # headway distance
        distance = compute_distance(ego_vehicle_loc, frontal_vehicle_loc)
        self.calculate_gap(distance)

        # Distance is computed from the center of the two cars,
        # use bounding boxes to calculate the actual distance
        distance = distance - max(
            frontal_vehicle.bounding_box.extent.y, frontal_vehicle.bounding_box.extent.x) - max(
            self.vehicle.bounding_box.extent.y, self.vehicle.bounding_box.extent.x)

        # safe control for car following todo: make the coefficient controllable
        if distance <= self._ego_speed / 3.6 * 0.01:
            print("emergency stop!")
            return 0, None

        target_speed, target_waypoint = self.platooning_following_manager(self.inter_gap)

        return target_speed, target_waypoint

    def run_step_cut_in_move2point(self):
        """
        The vehicle is trying to get to the move in point.

        Parameters
        ----------
        target_speed : float
            The target speed for ego vehile.
        target_waypoint : carla.waypoint
            The waypoint for ego vehile.
        next FSM state : enum
            The next finite state machine state.

        Returns
        -------
        """

        frontal_vehicle_manager, rear_vehicle_vm = self.v2x_manager.get_platoon_front_rear()
        frontal_vehicle = frontal_vehicle_manager.vehicle

        ego_vehicle_loc = self._ego_pos.location
        ego_vehicle_yaw = self._ego_pos.rotation.yaw

        distance, angle = cal_distance_angle(frontal_vehicle.get_location(),
                                             ego_vehicle_loc, ego_vehicle_yaw)

        # calculate the time gap with the frontal vehicle
        self.calculate_gap(distance)

        # if there is a obstacle blocking ahead, we just change to back joining mode todo: lane change not considered
        if self.hazard_flag:
            if rear_vehicle_vm:
                rear_vehicle_vm.v2x_manager.set_platoon_status(FSM.MAINTINING)
            # retrieve the last member in the platoon
            platoon_manager, _ = frontal_vehicle_manager.v2x_manager.get_platoon_manager()
            last_member = platoon_manager.vehicle_manager_list[-1]

            # set the last member as the frontal vehicle
            self.v2x_manager.set_platoon_front(last_member)
            self.v2x_manager.set_platoon_rear(None)
            print('switch to back joining!')
            # slow down to join back
            return (*super().run_step(self.max_speed / 2), FSM.BACK_JOINING)

        # the vehicle needs to warm up first. But if the platooning is in car following state, then we should ignore
        if self._ego_speed <= self.warm_up_speed and not frontal_vehicle_manager.agent.car_following_flag:
            print("warming up speed")
            return (*super().run_step(self.tailgate_speed), FSM.MOVE_TO_POINT)

        # if the ego vehicle is still too far away from the front vehicle
        if distance > self._ego_speed / 3.6 * (self.inter_gap + 0.5) and angle <= 80:
            print('trying to get the vehicle')
            return (*super().run_step(2.0 * get_speed(frontal_vehicle)), FSM.MOVE_TO_POINT)

        # if the ego vehicle is too close or exceed the frontal vehicle
        if distance < self._ego_speed / 3.6 * self.inter_gap / 1.5 or angle >= 70:
            print('too close, step back!')
            return (*super().run_step(0.9 * get_speed(frontal_vehicle)), FSM.MOVE_TO_POINT)

        # communicate to the rear vehicle for open gap if rear vehicle exists
        if not rear_vehicle_vm:
            return (*self.platooning_merge_management(frontal_vehicle_manager), FSM.JOINING)

        distance, angle = cal_distance_angle(rear_vehicle_vm.vehicle.get_location(),
                                             ego_vehicle_loc, ego_vehicle_yaw)

        # check whether the rear vehicle gives enough gap
        if distance < 1.0 * self.inter_gap / 2.0 * self._ego_speed / 3.6 \
                or angle <= 100 or rear_vehicle_vm.agent.current_gap < self.open_gap:
            # force the rear vehicle open gap for self
            print("too close to rear vehicle!")
            rear_vehicle_vm.v2x_manager.set_platoon_status(FSM.OPEN_GAP)
            return (*super().run_step(1.5 * get_speed(frontal_vehicle)), FSM.MOVE_TO_POINT)

        return (*self.platooning_merge_management(frontal_vehicle_manager), FSM.JOINING)

    def run_step_cut_in_joining(self):
        """
        Check if the vehicle has been joined successfully.
        """
        print("merging speed %d" % self._ego_speed)

        frontal_vehicle_manager, rear_vehicle_vm = self.v2x_manager.get_platoon_front_rear()

        frontal_vehicle = frontal_vehicle_manager.vehicle
        frontal_lane = self._map.get_waypoint(frontal_vehicle.get_location()).lane_id

        ego_vehicle_loc = self._ego_pos.location
        ego_vehicle_lane = self._map.get_waypoint(ego_vehicle_loc).lane_id
        ego_vehicle_yaw = self._ego_pos.rotation.yaw

        distance, angle = cal_distance_angle(frontal_vehicle.get_location(),
                                             ego_vehicle_loc, ego_vehicle_yaw)
        # calculate the time gap with the frontal vehicle
        self.calculate_gap(distance)

        if frontal_lane == ego_vehicle_lane and angle <= 5:
            print('merge finished')
            if rear_vehicle_vm:
                rear_vehicle_vm.v2x_manager.set_platoon_status(FSM.MAINTINING)
            return (*self.run_step_maintaining(), FSM.JOINING_FINISHED)

        return (*super().run_step(target_speed=get_speed(frontal_vehicle),
                                  collision_detector_enabled=False), FSM.JOINING)

    def run_step_open_gap(self):
        """
        Open gap for cut-in vehicle.
        """
        frontal_vehicle_manager, rear_vehicle_manager = self.v2x_manager.get_platoon_front_rear()

        # calculate the time gap under this state
        ego_vehicle_loc = self._ego_pos.location
        ego_vehicle_yaw = self._ego_pos.rotation.yaw

        distance, _ = cal_distance_angle(frontal_vehicle_manager.localizer.get_ego_pos().location,
                                         ego_vehicle_loc, ego_vehicle_yaw)
        self.calculate_gap(distance)

        # gradually open the gap TODO: Make this dynamic to map a linear relationship with speed
        if self.current_gap < self.open_gap:
            self.current_gap += 0.01
        print('cuurent gap is %f' % self.current_gap)
        target_speed, target_loc = self.platooning_following_manager(self.current_gap)

        return target_speed, target_loc

    def run_step_back_joining(self):
        """
        Back-joining Algorithm.

        Parameters
        ----------

        Returns
        -------
        control command : opencda object
            Control command for bacj joining.
        back join status : enum
            FSM back joining status.

        """
        frontal_vehicle_manager, _ = self.v2x_manager.get_platoon_front_rear()
        # reset lane change flag every step

        # get necessary information of the ego vehicle and target vehicle in the platooning
        frontal_vehicle = frontal_vehicle_manager.vehicle
        frontal_lane = self._map.get_waypoint(frontal_vehicle.get_location()).lane_id

        # retrieve the platooning's destination
        platooning_manager, _ = frontal_vehicle_manager.v2x_manager.get_platoon_manager()
        frontal_destination = platooning_manager.destination

        # retrieve ego vehicle info todo: remove this later
        ego_vehicle_loc = self._ego_pos.location
        ego_wpt = self._map.get_waypoint(ego_vehicle_loc)
        ego_vehicle_lane = ego_wpt.lane_id
        ego_vehicle_yaw = self._ego_pos.rotation.yaw

        distance, angle = cal_distance_angle(frontal_vehicle.get_location(),
                                             ego_vehicle_loc, ego_vehicle_yaw)

        # calculate the time gap with the frontal vehicle
        self.calculate_gap(distance)

        # 0. make sure the vehicle is behind the ego vehicle
        if angle >= 60 or distance < self._ego_speed / 3.6 * 0.5:
            self.overtake_allowed = False
            print("angle is too large, wait")
            return (*super().run_step(get_speed(frontal_vehicle) * 0.90, lane_change_allowed=False), FSM.BACK_JOINING)

        else:
            self.overtake_allowed = True

        # 1. make sure the speed is warmed up first. Also we don't want to reset destination during lane change
        if self._ego_speed < self.warm_up_speed or self.get_local_planner().lane_change:
            print('warm up speed')
            return (*super().run_step(self.tailgate_speed), FSM.BACK_JOINING)

        if not self.destination_changed:
            print('destination reset!!!!')
            self.destination_changed = True
            self.set_destination(ego_wpt.next(4.5)[0].transform.location, frontal_destination,
                                 clean=True, clean_history=True)

        # 2. check if there is any other vehicle blocking between ego and platooning
        def dist(v):
            return v.get_location().distance(ego_vehicle_loc)

        vehicle_blocking_status = False
        for vehicle in self.obstacle_vehicles:
            vehicle_blocking_status = vehicle_blocking_status or self._collision_check.is_in_range(self._ego_pos,
                                                                                frontal_vehicle,
                                                                                vehicle,
                                                                                self._map)

        # 3. if no other vehicle is blocking, the ego vehicle is in the same lane with the platooning
        # and it is close enough, then we regard the back joining finished
        if frontal_lane == ego_vehicle_lane \
                and not vehicle_blocking_status \
                and distance < 1.0 * self._ego_speed / 3.6:
            print('joining finished !')
            return (*self.run_step_maintaining(), FSM.JOINING_FINISHED)

        # 4. If vehicle is not blocked, make ego back to the frontal vehicle's lane
        if not vehicle_blocking_status:
            print('no vehicle is blocking!!!')
            if frontal_lane != ego_vehicle_lane:
                left_wpt = ego_wpt.next(max(1.2 * self._ego_speed / 3.6, 5))[0].get_left_lane()
                right_wpt = ego_wpt.next(max(1.2 * self._ego_speed / 3.6, 5))[0].get_right_lane()

                if not left_wpt and not right_wpt:
                    pass
                # if no right lane
                elif not right_wpt:
                    print('take left lane')
                    self.set_destination(left_wpt.transform.location, frontal_destination,
                                         clean=True, clean_history=True)
                # if no left lane available
                elif not left_wpt:
                    print('take right lane')
                    self.set_destination(right_wpt.transform.location, frontal_destination,
                                         clean=True, clean_history=True)
                # check which lane is closer to the platooning
                elif abs(left_wpt.lane_id - frontal_lane) < abs(right_wpt.lane_id - frontal_lane):
                    print('take left lane')
                    self.set_destination(left_wpt.transform.location, frontal_destination,
                                         clean=True, clean_history=True)
                else:
                    print('take right lane')
                    self.set_destination(right_wpt.transform.location, frontal_destination,
                                         clean=True, clean_history=True)

        return (*super().run_step(self.tailgate_speed), FSM.BACK_JOINING)

    def run_step_front_joining(self):
        """
        Front-joining algorithm.

        Parameters
        ----------

        Returns
        -------
        control command : opencda object
            Control command for bacj joining.
        back join status : enum
            FSM back joining status.
        """
        _, rear_vehicle_manager = self.v2x_manager.get_platoon_front_rear()
        # get necessary information of the ego vehicle and target vehicle in the platooning
        rear_vehicle = rear_vehicle_manager.vehicle
        rear_lane = self._map.get_waypoint(rear_vehicle.get_location()).lane_id

        # retrieve the platooning's destination
        platooning_manager, _ = rear_vehicle_manager.v2x_manager.get_platoon_manager()
        rear_destination = platooning_manager.destination

        # retrieve ego vehicle info
        ego_vehicle_loc = self._ego_pos.location
        ego_wpt = self._map.get_waypoint(ego_vehicle_loc)
        ego_vehicle_lane = ego_wpt.lane_id
        ego_vehicle_yaw = self._ego_pos.rotation.yaw

        distance, angle = cal_distance_angle(rear_vehicle.get_location(),
                                             ego_vehicle_loc, ego_vehicle_yaw)

        # if there is a vehicle blocking between, then abandon this joining
        def dist(v):
            return v.get_location().distance(ego_vehicle_loc)

        vehicle_blocking_status = False
        for vehicle in self.obstacle_vehicles:
            vehicle_blocking_status = vehicle_blocking_status or self._collision_check.is_in_range(self._ego_pos,
                                                                                rear_vehicle,
                                                                                vehicle,
                                                                                self._map)
        # if vehicle blocking between ego and platooning, then abandon this joining
        if vehicle_blocking_status:
            print('abandon front joining')
            return (*super().run_step(self.max_speed - self.speed_lim_dist), FSM.ABONDON)

        # if the ego vehilce is already behind the platooning
        if angle <= 90:
            print('transition to other joining')
            return (*super().run_step(self.tailgate_speed), FSM.SEARCHING)

        # if vehicle is already in the same lane with the platooning
        if ego_vehicle_lane == rear_lane:
            # after in the same lane, no overtake is allowed anymore
            self.overtake_allowed = False

            if not self.destination_changed:
                self.set_destination(ego_wpt.next(4.5)[0].transform.location, rear_destination)
            if distance < self._ego_speed / 3.6 * self.inter_gap * 1.5:
                print('joining finished')
                return (*super().run_step(self.max_speed - self.speed_lim_dist), FSM.JOINING_FINISHED)
            else:
                return (*super().run_step(get_speed(rear_vehicle) * 0.95), FSM.FRONT_JOINING)

        # if the ego is too close to the platooning or speed is too slow
        if distance < self._ego_speed / 3.6 * self.inter_gap \
                or self._ego_speed < self.warm_up_speed \
                or angle <= 90 or get_speed(rear_vehicle) > self._ego_speed \
                or self.get_local_planner().lane_change:
            print('need to speed up before change lane')
            return (*super().run_step(self.tailgate_speed), FSM.FRONT_JOINING)

        # set destination same as platooning
        if not self.destination_changed:
            print('destination reset!!!!')
            self.destination_changed = True
            self.set_destination(ego_wpt.next(4.5)[0].transform.location, rear_destination, clean=True,
                                 clean_history=True)

        # check which lane is closer to operate lane change
        left_wpt = ego_wpt.next(max(1.2 * self._ego_speed / 3.6, 5))[0].get_left_lane()
        right_wpt = ego_wpt.next(max(1.2 * self._ego_speed / 3.6, 5))[0].get_right_lane()
        # check which lane is closer to the platooning
        if not left_wpt and not right_wpt:
            pass
        # check which lane is closer to the platooning
        elif not right_wpt or abs(left_wpt.lane_id - rear_lane) < abs(right_wpt.lane_id - rear_lane):
            print('take left lane')
            self.set_destination(left_wpt.transform.location, rear_destination, clean=True)
        else:
            print('take right lane')
            self.set_destination(right_wpt.transform.location, rear_destination, clean=True)

        return (*super().run_step(self.tailgate_speed), FSM.FRONT_JOINING)

ext: py

reason: import numpy


text:
import matplotlib as mpl
import matplotlib.pyplot as plt
mpl.use('Qt5Agg')
import tensorflow as tf
import numpy as np
import keras
import PIL

import pathlib

import keras.backend as K
from tensorflow.python.framework import ops

from scipy import ndimage
import copy

from collections import OrderedDict
from PyQt5.QtWidgets import QWidget, QCheckBox, QApplication
from PyQt5.QtCore import Qt

import sys

sys.path.append("DeepDanbooru")

import AnimeFaceNotebooks.DeepDanbooru.deepdanbooru as dd

import pickle
import glob
import time

from ipywidgets import interact, interactive, fixed, interact_manual
import ipywidgets as widgets

from IPython.display import clear_output, display
import io

# load model

#project_path = r"C:\Users\david\Desktop\Quintuplets2\AnimeFaceNotebooks\deepdanbooru_model"
project_path = str(pathlib.Path().absolute()) + r"\AnimeFaceNotebooks\deepdanbooru_model"
#print(projectPath)
#if(project_path == projectPath):
    #print("same")

classif_model = dd.project.load_model_from_project(project_path)
all_tags = dd.project.load_tags_from_project(project_path)
model_width = classif_model.input_shape[2]
model_height = classif_model.input_shape[1]
all_tags = np.array(all_tags)

# preprocessing of read the image
def read_image(image_path):
    image = dd.data.load_image_for_evaluate(image_path, width=model_width, height=model_height)
    image_fullres = PIL.Image.open(image_path).convert('RGB')
    return image_fullres, image

# Functions for tag estimation
##
def decode_tags(image_arrays, threshold=0.1):
    """
    Decode tags using DeepDanbooru, threshold them and return them
    as a sorted dict (highest likelihood first)

    Will return a list if a list is passed.
    """
    if len(np.array(image_arrays).shape) == 4:
        return_list = True
    else:
        image_arrays = np.array([image_arrays])
        return_list = False

    # Decode
    results = classif_model.predict(image_arrays).reshape(-1, all_tags.shape[0])

    # Threshold and sort
    result_list = []
    for result_set in results:
        result_tags = {}
        for i in range(len(all_tags)):
            if result_set[i] > threshold:
                result_tags[all_tags[i]] = result_set[i]

        sorted_tags = reversed(sorted(result_tags.keys(), key=lambda x: result_tags[x]))
        sorted_results = OrderedDict()
        for tag in sorted_tags:
            sorted_results[tag] = result_tags[tag]
        result_list.append(sorted_results)

    if return_list:
        return result_list
    else:
        return result_list[0]

def show_tags(result_tags):
    """
    Tag list prettyprinter
    """
    fo = open("output.txt", "w+")
    for tag in result_tags.keys():
        str='{0: <32} {1:04f}'.format(tag, result_tags[tag])
        fo.write(str)
        fo.write("\n")
    return fo

def get_likelihood(image_array, tags):
    """
    Decode and return mean likelihood of given tags
    """
    results = decode_tags(image_array, threshold=0.0)
    sum_likelihood = 0.0
    for tag in tags:
        sum_likelihood += results[tag]
    return sum_likelihood / len(tags)

'''
image_fullres, image = read_image("../gcat.jpg")
plt.imshow(image_fullres)
tags = decode_tags(image)
show_tags(tags)
'''
# Gradcam with batching
##
def gradcam(image_array, class_sets, batch_size=5, verbose=False):
    """
    Get gradient at inputs that maximizes the given class(es)
    """
    results = []
    for batch_start in range(0, len(class_sets), batch_size):
        if verbose:
            print("Batch: {} to {} of {}".format(batch_start, batch_start + batch_size, len(class_sets)))
        batch = class_sets[batch_start:batch_start + batch_size]

        batch_size_real = len(batch)
        image_array_feed = np.array([image_array] * batch_size_real)
        image_feed = tf.Variable(image_array_feed)

        batch_nhot_list = []
        for classes in batch:
            if not isinstance(classes, list):
                classes = [classes]

            class_nhot = np.zeros(len(all_tags))
            for gradcam_class in classes:
                class_idx = list(all_tags).index(gradcam_class)
                class_nhot[class_idx] = 1.0
            batch_nhot_list.append(class_nhot)

        with tf.GradientTape() as tape:
            output = classif_model(image_feed)
            test = tf.multiply(batch_nhot_list, output)
            gradcam_loss = tf.reduce_sum(tf.multiply(batch_nhot_list, output))
            results.extend(list(np.array(tape.gradient(gradcam_loss, image_feed))))
    return np.array(results)

    # Functions for activation map processing

##
def norm_clip_grads(grads):
    """
    Normalize gradcam gradients by clipping them to between the 1st and
    the 99th percentile, taking the absolute value and then mapping that
    range to 0 -> 1
    """
    upper_quantile = np.quantile(grads, 0.99)
    lower_quantile = np.quantile(grads, 0.01)
    clipped_grads = np.abs(np.clip(grads, lower_quantile, upper_quantile))
    return clipped_grads / np.max(clipped_grads)

def postprocess_grads(grads, filter_width=10):
    """
    Normalize gradients, take the mean and apply
    a median filter
    """
    gradcam_map = np.mean(norm_clip_grads(grads), axis=-1)
    return ndimage.median_filter(gradcam_map, filter_width)

def threshold_map(gradcam_map, all_maps, filter_width=10):
    """
    Refine gradcam map by applying a threshold mask that zeroes out
    areas where the map is not bigger than the average of all maps
    (i.e. remove areas that are only "generally interesting" instead
    of interesting for this tag specifically)
    """
    threshold_mask = gradcam_map > np.mean(all_maps, axis=0)
    return ndimage.median_filter(threshold_mask, filter_width) * gradcam_map

def mask_for_tags(tags, gradcam_maps_processed, clip_value=0.3, regu_value=0.0, median_filter_width=10, anti_tags=None):
    """
    Given tag(s) and a set of (normalized) gradcam maps, generate a mask of pixels
    that contribute to the threshold_map image being assigned those tags

    Maps are combined multiplicatively, anti-tags (if specified) are used inversely
    """
    if not isinstance(tags, list):
        tags = [tags]

    tags_map = np.ones(gradcam_maps_processed[tags[0]].shape)
    for tag in tags:
        tag_map = threshold_map(gradcam_maps_processed[tag], list(gradcam_maps_processed.values()))
        tag_map = ndimage.median_filter(tag_map, median_filter_width)
        tag_map /= np.max(tag_map)
        tags_map *= tag_map

    if not anti_tags is None:
        for tag in anti_tags:
            if tag in gradcam_maps_processed:
                tag_map = threshold_map(gradcam_maps_processed[tag], list(gradcam_maps_processed.values()))
                tag_map = ndimage.median_filter(tag_map, median_filter_width)
                tag_map /= np.max(tag_map)
                tags_map *= (1.0 - tag_map)

    tags_map /= np.max(tags_map)
    tags_map = np.maximum(tags_map, clip_value)
    tags_map += regu_value
    tags_map /= np.max(tags_map)

    return tags_map

def show_mask(image, mask):
    """
    Show a nice plot of an image, a mask, and the masked image
    """
    plt.figure(figsize=(16, 5))
    plt.subplot(1, 3, 1)
    plt.imshow(image)
    plt.xlabel("Input image")

    plt.subplot(1, 3, 2)
    plt.imshow(mask, vmin=0.0, vmax=1.0)
    plt.xlabel("Activation mask")

    plt.subplot(1, 3, 3)
    plt.imshow(image * np.repeat(mask, 3).reshape(image.shape))
    plt.xlabel("Masked image")

def plot_tags(image_arr, gradcam_maps_processed, tags, anti_tags=None):
    """
    Get mask for tag and plot
    """
    mask = mask_for_tags(tags, gradcam_maps_processed, 0.2, anti_tags=anti_tags)
    show_mask(image_arr, mask)
    title = "Mask for tags: '{}'".format(" ".join(tags))
    if not anti_tags is None:
        title += " excluding '{}'".format(" ".join(anti_tags))
    plt.suptitle(title)

# Postprocessing and masking tools
##
def get_mask_image(gradcam_maps_processed, tags_to_mask, tags_to_keep=None, mask_thresh=0.2, inverse=False,
                   dilation_scale=1.0):
    """
    Find a binary tag mask for one or more tags by thresholding and
    dilation/erosion
    """
    mask = np.zeros(gradcam_maps_processed[list(gradcam_maps_processed.keys())[0]].shape)
    if not isinstance(tags_to_mask, list):
        tags_to_mask = [tags_to_mask]

    tag_set_real = []
    for tag in tags_to_mask:
        if tag in gradcam_maps_processed:
            tag_set_real.append(tag)

    if len(tag_set_real) > 0:
        temp_mask = mask_for_tags(tag_set_real, gradcam_maps_processed, 0.0, anti_tags=tags_to_keep) < mask_thresh
        mask += 1.0 - temp_mask

    mask = 1.0 - np.clip(mask, 0.0, 1.0)
    if inverse:
        mask = 1.0 - mask

    # Do some hole / noise removal and make the mask a tad larger
    mask = ndimage.morphology.binary_dilation(mask, iterations=int(5 * dilation_scale), border_value=1)
    mask = ndimage.morphology.binary_erosion(mask, iterations=int(15 * dilation_scale), border_value=1) * 1.0

    return mask

def mask_image_array(image_array, mask):
    """
    Apply a mask to an image array
    """
    return image_array * mask.reshape(mask.shape[0], mask.shape[1], 1).repeat(3, axis=-1)

def find_tuned_mask_image(image_array, gradcam_maps_processed, tags_to_mask, tags_to_keep=None, threshold_factor=0.25,
                          maximum_likelihood=1.0):
    """
    Get a mask image and use the classifier to automatically tune the
    threshold for masking by looking for a setting for which the
    likelihood is under a certain value.

    Returns the mask and the determined threshold
    """
    # Figure out at what threshold we should stop
    initial_likelihood = get_likelihood(image_array, tags_to_mask)
    if initial_likelihood < 0.01:
        return np.ones(image_array.shape[0:2]), 1.0
    break_likelihood = min(maximum_likelihood, initial_likelihood * threshold_factor)

    # Now run until we go past that and return
    for i in range(0, 70):
        threshold = (70 + 1) * 0.01 - 0.01 * i
        mask = get_mask_image(gradcam_maps_processed, tags_to_mask, tags_to_keep, mask_thresh=threshold)
        likelihood = get_likelihood(mask_image_array(image_array, mask), tags_to_mask)
        if likelihood < break_likelihood:
            break
    return mask, threshold

# Full-image masking tools (with pixelation for fun. christian anime websites: call me)
# All of these are uneccesarily complicated and could be simplified a lot, probably
##
def mask_fullres_image(image_fullres, image_array, gradcam_maps_processed, tags_to_mask, tags_to_keep=None,
                       inverse=False, pixelate=True, dilation_scale=1.0):
    """
    Find a tag mask and apply it to a full-resolution image (either blacking
    out or pixelating regions)

    This is a bit slow, but should be fairly automatic
    """
    # Get a combined (additive) mask for these tags
    mask = np.zeros(gradcam_maps_processed[list(gradcam_maps_processed.keys())[0]].shape)
    for tag_set in tags_to_mask:
        if not isinstance(tag_set, list):
            tag_set = [tag_set]
        tmp_mask, _ = find_tuned_mask_image(image_array, gradcam_maps_processed, tag_set, tags_to_keep)
        mask += 1.0 - tmp_mask
    mask = np.clip(mask, 0.0, 1.0)
    if not inverse:
        mask = 1.0 - mask

    # Scale mask to image size
    mask_image = np.repeat(mask.reshape(mask.shape[0], mask.shape[1], 1), 3, axis=-1) * 255.0
    mask_image = PIL.Image.fromarray(mask_image.astype("uint8"))
    mask_image = mask_image.resize(image_fullres.size, resample=PIL.Image.LANCZOS)

    # Gauss filter to make everything look nice and soft
    mask = np.array(mask_image)[:, :, 0] / 255.0
    mask = ndimage.gaussian_filter(mask, 3)

    # Combined input and pixelated using mask
    image_fullres_arr = np.array(image_fullres)
    masked_image = np.array(image_fullres_arr * np.repeat(mask, 3).reshape(image_fullres_arr.shape))

    # Create a pixelated version of the image if desired
    if pixelate == True:
        image_fullres_pixeled = image_fullres.resize(np.array(image_fullres.size) // 15).resize(image_fullres.size,
                                                                                resample=PIL.Image.NEAREST)
        pixeled_arr = np.array(image_fullres_pixeled)

        # Combine with base masked image
        pixeled_image = np.array(pixeled_arr * np.repeat(1.0 - mask, 3).reshape(image_fullres_arr.shape))
        masked_image = masked_image + pixeled_image

    # Return masked image, as a PIL image
    return PIL.Image.fromarray(masked_image.astype("uint8"))

def mask_fullres_image_simple(image_fullres, image_array, gradcam_maps_processed, tags_to_mask, tags_to_keep=None,
                              inverse=False, pixelate=True, dilation_scale=1.5, threshold=0.2):
    """
    Find a tag mask and apply it to a full-resolution image (either blacking
    out or pixelating regions)

    Faster and worse
    """

    # Get a combined (additive) mask for these tags
    mask = np.zeros(gradcam_maps_processed[list(gradcam_maps_processed.keys())[0]].shape)
    for tag_set in tags_to_mask:
        if not isinstance(tag_set, list):
            tag_set = [tag_set]
        tmp_mask = get_mask_image(gradcam_maps_processed, tag_set, tags_to_keep, mask_thresh=threshold,
                                  dilation_scale=dilation_scale)
        mask += 1.0 - tmp_mask
    mask = np.clip(mask, 0.0, 1.0)
    if not inverse:
        mask = 1.0 - mask

    # Scale mask to image size
    mask_image = np.repeat(mask.reshape(mask.shape[0], mask.shape[1], 1), 3, axis=-1) * 255.0
    mask_image = PIL.Image.fromarray(mask_image.astype("uint8"))
    mask_image = mask_image.resize(image_fullres.size, resample=PIL.Image.LANCZOS)

    # Gauss filter to make everything look nice and soft
    mask = np.array(mask_image)[:, :, 0] / 255.0
    mask = ndimage.gaussian_filter(mask, 3)

    # Combined input and pixelated using mask
    image_fullres_arr = np.array(image_fullres)
    masked_image = np.array(image_fullres_arr * np.repeat(mask, 3).reshape(image_fullres_arr.shape))

    # Create a pixelated version of the image if desired
    if pixelate == True:
        image_fullres_pixeled = image_fullres.resize(np.array(image_fullres.size) // 15).resize(image_fullres.size,
                                                                                resample=PIL.Image.NEAREST)
        pixeled_arr = np.array(image_fullres_pixeled)

        # Combine with base masked image
        pixeled_image = np.array(pixeled_arr * np.repeat(1.0 - mask, 3).reshape(image_fullres_arr.shape))
        masked_image = masked_image + pixeled_image

    # Return masked image, as a PIL image
    return PIL.Image.fromarray(masked_image.astype("uint8"))

image = None
image_fullres = None
tags = None
gradcam_maps = None
gradcam_maps_processed = None
tag_maps = None
censor_button = None

# Interactive tagging + gradcam visualization (by censoring images)
# Drag slider towards 1 to pixelate, towards -1 to exclude
# View Output Fullscreen (on the cell below this one) recommended
##
tag_widgets = None

def upload_image(b):
    file_name = list(b["new"].keys())[0]
    image = PIL.Image.open(io.BytesIO(b["new"][file_name]["content"]))
    image.save("input.png")
    load_image("input.png")

def load_image(path):
    # Load an image
    global image
    global image_fullres
    global tags
    global gradcam_maps
    global gradcam_maps_processed
    global tag_maps
    global tag_widgets
    global censor_button

    image_fullres, image = read_image(path)

    print("Tagging...")
    tags = decode_tags(image)

    print("Calculating gradcam maps...")
    gradcam_maps = dict(zip(list(tags.keys()), gradcam(image, list(tags.keys()), batch_size=5, verbose=True)))

    print("Postprocessing gradcam maps...")
    gradcam_maps_processed = dict(zip(gradcam_maps.keys(), list(map(postprocess_grads, gradcam_maps.values()))))

    print("Calculating tag maps...")
    tag_maps = {}
    for tag in gradcam_maps_processed:
        tag_map = threshold_map(gradcam_maps_processed[tag], list(gradcam_maps_processed.values()))
        tag_map = ndimage.median_filter(tag_map, 10)
        tag_map /= np.max(tag_map)
        tag_maps[tag] = tag_map

    tag_widgets = {}
    for tag in tags:
        tag_widgets[tag] = QWidget.FloatSlider(description=tag, min=-0.99, max=0.99, step=0.01, value=0.0,
                                               continuous_update=False, readout=True)
    all_widgets = list(tag_widgets.values())

    censor_button = QWidget.ToggleButton(value=False, description='Pixelate')

    reset_sliders = QWidget.Button(description="Reset")
    reset_sliders.on_click(reset)

    #upload = QWidget.FileUpload(
    #    description="Upload image",
    #    accept=".png,.jpg,.jpeg"
    #)
    #upload.observe(upload_image, names="value")

    ui = QWidget.Box(all_widgets + [reset_sliders, censor_button])

    ui.layout.flex_flow = 'row wrap'
    ui.layout.display = 'inline-flex'
    tag_widgets["pixelate"] = censor_button

    out = QWidget.interactive_output(censor_and_display, tag_widgets)
    clear_output(wait=True)
    return display(ui, out)

def censor_and_display(pixelate, **kwargs):
    mask = np.zeros(gradcam_maps_processed[list(gradcam_maps_processed.keys())[0]].shape)
    for tag in kwargs:
        if kwargs[tag] > 0.0:
            mask += (tag_maps[tag] > (1.0 - kwargs[tag]))
        if kwargs[tag] < 0.0:
            mask -= (tag_maps[tag] > (1.0 + kwargs[tag]))

    if pixelate:
        mask = 1.0 - (mask > 0.0)
    else:
        mask = mask > 0.0

    mask = ndimage.morphology.binary_dilation(mask, iterations=int(5 * 1.5), border_value=1)
    mask = ndimage.morphology.binary_erosion(mask, iterations=int(15 * 1.5), border_value=1)

    mask_image = np.repeat(mask.reshape(mask.shape[0], mask.shape[1], 1), 3, axis=-1) * 255.0
    mask_image = PIL.Image.fromarray(mask_image.astype("uint8"))
    mask_image = mask_image.resize(image_fullres.size, resample=PIL.Image.LANCZOS)

    mask = np.array(mask_image)[:, :, 0] / 255.0
    mask = ndimage.gaussian_filter(mask, 3)

    #if not pixelate:
      #  mask = mask * 0.8 + 0.2

    image_fullres_arr = np.array(image_fullres)
    masked_image = np.array(image_fullres_arr * np.repeat(mask, 3).reshape(image_fullres_arr.shape))
    if pixelate:
        image_fullres_pixeled = image_fullres.resize(np.array(image_fullres.size) // 15).resize(image_fullres.size,
                                                                                resample=PIL.Image.NEAREST)
        pixeled_arr = np.array(image_fullres_pixeled)

        pixeled_image = np.array(pixeled_arr * np.repeat(1.0 - mask, 3).reshape(image_fullres_arr.shape))
        masked_image = masked_image + pixeled_image

    censored = PIL.Image.fromarray(masked_image.astype("uint8"))
    display(censored)

def reset(b):
    for widget in tag_widgets:
        if not widget == "pixelate":
            tag_widgets[widget].value = 0.0

def load_input(b, image):
    load_image("input.png")

max_stars_repo_name: davidiswhat/AnimeIllustrationRecognizer

max_issues_repo_name: davidiswhat/AnimeIllustrationRecognizer

max_forks_repo_name: davidiswhat/AnimeIllustrationRecognizer

reason: import numpy,from scipy


text:
from abc import abstractmethod

import os
import numpy as np
import torch
from numpy import inf
from utils import MetricTracker, inf_loop

class Trainer:
    """
    Base class for all trainers
    """
    def __init__(self, model, criterion, metric_ftns, optimizer, config,
                 device,
                 train_data_loader,
                 valid_data_loader=None,
                 test_data_loader=None,
                 lr_scheduler=None,
                 checkpoint_dir=None):
        self.config = config
        self.logger = config.get_logger('trainer', config['trainer']['verbosity'])
        self.device = device

        self.model = model.to(device)
        self.criterion = criterion
        self.metric_ftns = metric_ftns
        self.optimizer = optimizer
        self.lr_scheduler = lr_scheduler

        cfg_trainer = config['trainer']
        self.epochs = cfg_trainer['epochs']
        self.save_period = cfg_trainer['save_period']
        self.monitor = cfg_trainer.get('monitor', 'off')

        # configuration to monitor model performance and save best
        if self.monitor == 'off':
            self.mnt_mode = 'off'
            self.mnt_best = 0
        else:
            self.mnt_mode, self.mnt_metric = self.monitor.split()
            assert self.mnt_mode in ['min', 'max']

            self.mnt_best = inf if self.mnt_mode == 'min' else -inf
            self.early_stop = cfg_trainer.get('early_stop', inf)
            if self.early_stop <= 0:
                self.early_stop = inf

        self.start_epoch = 1
        self.checkpoint_dir = checkpoint_dir
        if checkpoint_dir is None:
            self.checkpoint_dir = config.save_dir
        if not os.path.exists(self.checkpoint_dir):
            os.mkdir(self.checkpoint_dir)
        else:
            for filename in os.listdir(self.checkpoint_dir):
                os.remove(os.path.join(self.checkpoint_dir, filename))

        self.train_data_loader = train_data_loader
        self.valid_data_loader = valid_data_loader
        self.test_data_loader = test_data_loader
        self.do_validation = self.valid_data_loader is not None
        self.log_step = int(np.sqrt(train_data_loader.batch_size))
        self.len_epoch = len(self.train_data_loader)

        self.train_metrics = MetricTracker('loss', *[m.__name__ for m in self.metric_ftns])
        self.valid_metrics = MetricTracker('loss', *[m.__name__ for m in self.metric_ftns])

    def _train_epoch(self, epoch):
        """
        Training logic for an epoch

        :param epoch: Integer, current training epoch.
        :return: A log that contains average loss and metric in this epoch.
        """
        self.model.train()
        self.train_metrics.reset()
        for batch_idx, (data, target) in enumerate(self.train_data_loader):
            data, target = data.to(self.device), target.to(self.device)

            self.optimizer.zero_grad()
            output = self.model(data)
            loss = self.criterion(output, target)
            loss.backward()
            self.optimizer.step()

            self.train_metrics.update('loss', loss.item())
            for met in self.metric_ftns:
                self.train_metrics.update(met.__name__, met(output, target))

            if batch_idx == self.len_epoch:
                break
        log = self.train_metrics.result()

        if self.do_validation:
            val_log = self._valid_epoch(epoch)
            log.update(**{'val_'+k : v for k, v in val_log.items()})

        if self.lr_scheduler is not None:
            self.lr_scheduler.step()
        return log

    def _valid_epoch(self, epoch):
        """
        Validate after training an epoch

        :param epoch: Integer, current training epoch.
        :return: A log that contains information about validation
        """
        self.model.eval()
        self.valid_metrics.reset()
        with torch.no_grad():
            for batch_idx, (data, target) in enumerate(self.valid_data_loader):
                data, target = data.to(self.device), target.to(self.device)

                output = self.model(data)
                loss = self.criterion(output, target)

                self.valid_metrics.update('loss', loss.item())
                for met in self.metric_ftns:
                    self.valid_metrics.update(met.__name__, met(output, target))

        return self.valid_metrics.result()

    def _progress(self, batch_idx):
        base = '[{}/{} ({:.0f}%)]'
        if hasattr(self.train_data_loader, 'n_samples'):
            current = batch_idx * self.train_data_loader.batch_size
            total = self.train_data_loader.n_samples
        else:
            current = batch_idx
            total = self.len_epoch
        return base.format(current, total, 100.0 * current / total)

    def train(self):
        """
        Full training logic
        """
        not_improved_count = 0
        for epoch in range(self.start_epoch, self.epochs + 1):
            result = self._train_epoch(epoch)

            # save logged informations into log dict
            log = {'epoch': epoch}
            log.update(result)

            # print logged informations to the screen
            for key, value in log.items():
                self.logger.info('    {:15s}: {}'.format(str(key), value))

            # evaluate model performance according to configured metric, save best checkpoint as model_best
            best = False
            if self.mnt_mode != 'off':
                try:
                    # check whether model performance improved or not, according to specified metric(mnt_metric)
                    improved = (self.mnt_mode == 'min' and log[self.mnt_metric] <= self.mnt_best) or \
                               (self.mnt_mode == 'max' and log[self.mnt_metric] >= self.mnt_best)
                except KeyError:
                    self.logger.warning("Warning: Metric '{}' is not found. "
                                        "Model performance monitoring is disabled.".format(self.mnt_metric))
                    self.mnt_mode = 'off'
                    improved = False

                if improved:
                    self.mnt_best = log[self.mnt_metric]
                    not_improved_count = 0
                    best = True
                else:
                    not_improved_count += 1

                if not_improved_count > self.early_stop:
                    self.logger.info("Validation performance didn\'t improve for {} epochs. "
                                     "Training stops.".format(self.early_stop))
                    break

            if best:
                self._save_checkpoint(epoch)

    def test(self):
        best_path = os.path.join(self.checkpoint_dir, f'model_best.pth')
        if os.path.exists(best_path):
            self.model.load_state_dict(torch.load(best_path, map_location=self.device)["state_dict"])

        self.model.eval()

        total_loss = 0.0
        total_metrics = torch.zeros(len(self.metric_ftns))

        with torch.no_grad():
            for i, (data, target) in enumerate(self.test_data_loader):
                data, target = data.to(self.device), target.to(self.device)
                output = self.model(data)

                # computing loss, metrics on test set
                loss = self.criterion(output, target)
                batch_size = data.shape[0]
                total_loss += loss.item() * batch_size
                for i, metric in enumerate(self.metric_ftns):
                    total_metrics[i] += metric(output, target) * batch_size

        n_samples = len(self.test_data_loader.sampler)
        log = {'loss': total_loss / n_samples}
        log.update({
            met.__name__: total_metrics[i].item() / n_samples for i, met in enumerate(self.metric_ftns)
        })
        self.logger.info(log)

        return log["accuracy"]

    def _save_checkpoint(self, epoch):
        """
        Saving checkpoints

        :param epoch: current epoch number
        :param log: logging information of the epoch
        :param save_best: if True, rename the saved checkpoint to 'model_best.pth'
        """
        arch = type(self.model).__name__
        state = {
            'arch': arch,
            'epoch': epoch,
            'state_dict': self.model.state_dict(),
        }
        self.logger.info("Best checkpoint in epoch {}".format(epoch))

        best_path = os.path.join(self.checkpoint_dir, 'model_best.pth')
        torch.save(state, best_path)
        self.logger.info("Saving current best: model_best.pth ...")

    def load_best_model(self):
        # Load the best model then test
        arch = type(self.model).__name__
        best_path = os.path.join(self.checkpoint_dir, 'model_best.pth')
        state_dict  = torch.load(best_path, map_location=self.device)["state_dict"]
        self.model.load_state_dict(state_dict)
        return state_dict

    # not in use
    def _resume_checkpoint(self, resume_path):
        """
        Resume from saved checkpoints

        :param resume_path: Checkpoint path to be resumed
        """
        resume_path = str(resume_path)
        self.logger.info("Loading checkpoint: {} ...".format(resume_path))
        checkpoint = torch.load(resume_path)
        self.start_epoch = checkpoint['epoch'] + 1
        self.mnt_best = checkpoint['monitor_best']

        # load architecture params from checkpoint.
        if checkpoint['config']['arch'] != self.config['arch']:
            self.logger.warning("Warning: Architecture configuration given in config file is different from that of "
                                "checkpoint. This may yield an exception while state_dict is being loaded.")
        self.model.load_state_dict(checkpoint['state_dict'])

        # load optimizer state from checkpoint only when optimizer type is not changed.
        if checkpoint['config']['optimizer']['type'] != self.config['optimizer']['type']:
            self.logger.warning("Warning: Optimizer type given in config file is different from that of checkpoint. "
                                "Optimizer parameters not being resumed.")
        else:
            self.optimizer.load_state_dict(checkpoint['optimizer'])

        self.logger.info("Checkpoint loaded. Resume training from epoch {}".format(self.start_epoch))

reason: import numpy,from numpy


text:
"""
common.py.

This module contains common functions used in other modules.
"""
from datetime import datetime

import numpy as np
import pandas as pd

def df_info(df: pd.DataFrame, clean: bool = False) -> pd.DataFrame:
    """
    Analyze DataFrame column definitions.

    Parameters
    ----------
    df : pandas.DataFrame
        Input DataFrame for analysis.
    clean : bool
        Whether to perform initial cleaning of DataFrame.

    Returns
    ------
    pandas.DataFrame
        Object with DataFrame analysis.
    """
    start = datetime.now()
    if clean:
        df.fillna(np.nan, inplace=True)  # replace None with np.nan
        df = df.applymap(
            lambda x: x.strip() if isinstance(x, str) else x
        )  # trim spaces
        df = df.replace("", np.nan)
    dmem = df.memory_usage(deep=True).to_dict()
    dmem.pop("Index", None)
    dnull = dict((v, df[v].isna().any()) for v in df.columns.values)
    dcnt = dict((v, len(df[v].dropna())) for v in df.columns.values)
    dunicode = dict(
        (
            v,
            df[v]
            .astype("object")
            .apply(
                lambda r: len(r) != len(r.encode())
                if isinstance(r, str) == str
                else False
            )
            .any(),
        )
        for v in df.columns.values
        if (df[v].notna().any() and df[v].dtypes.kind == "O")
    )
    dmin = dict(
        (v, df[v].min())
        for v in df.columns.values
        if (df[v].notna().any() and df[v].dtypes.kind not in ["O", "f"])
    )
    dmax = dict(
        (v, df[v].max())
        for v in df.columns.values
        if (df[v].notna().any() and df[v].dtypes.kind not in ["O", "f"])
    )
    dmin.update(
        dict(
            (v, int(df[v].min()))
            for v in df.columns.values
            if (df[v].notna().any() and df[v].dtypes.kind == "f")
        )
    )
    dmax.update(
        dict(
            (v, int(df[v].max()))
            for v in df.columns.values
            if (df[v].notna().any() and df[v].dtypes.kind == "f")
        )
    )
    dmin.update(
        dict(
            (
                v,
                df[v]
                .apply(lambda r: len(str(r)) if r is not np.nan else np.nan)
                .astype("Int32")
                .min(),
            )
            for v in df.columns.values
            if (df[v].notna().any() and df[v].dtypes.kind == "O")
        )
    )
    dmax.update(
        dict(
            (
                v,
                df[v]
                .apply(lambda r: len(str(r)) if r is not np.nan else np.nan)
                .astype("Int32")
                .max(),
            )
            for v in df.columns.values
            if (df[v].notna().any() and df[v].dtypes.kind == "O")
        )
    )
    dmode = dict(
        (v, df[v].mode()[0])
        for v in df.columns.values
        if (df[v].notna().any() and not df[v].is_unique)
    )
    dmodecnt = dict(
        (v, len(df[df[v] == df[v].mode()[0]]))
        for v in df.columns.values
        if df[v].notna().any()
    )
    dunique = dict(
        (v, len(df[v].unique())) for v in df.columns.values if df[v].notna().any()
    )
    desc = pd.DataFrame({"types": df.dtypes})
    desc["size"] = pd.Series(dmem)
    desc["isnull"] = pd.Series(dnull)
    desc["count"] = pd.Series(dcnt)
    desc["isunicode"] = pd.Series(dunicode)
    desc["min"] = pd.Series(dmin)
    desc["max"] = pd.Series(dmax)
    desc["mode"] = pd.Series(dmode)
    desc["mode_cnt"] = pd.Series(dmodecnt).astype("Int64")
    desc["uniq_cnt"] = pd.Series(dunique).astype("Int64")
    print(f"\033[92mElapsed time: {datetime.now() - start}\033[0m")
    return desc

def tidy_split(df, column, sep="|", keep=False):
    """
    Split the values of a column to rows.

    Expands column values so the new DataFrame has one split value per row.
    Filters rows where the column is missing.

    Parameters
    ----------
    df : pandas.DataFrame
        dataframe with the column to split and expand.
    column : str
        the column to split and expand
    sep : str
        the string used to split the column's values
    keep : bool
        whether to retain the presplit value as it's own row

    Returns
    -------
    pandas.DataFrame
        Returns a dataframe with the same columns as `df`.
    """
    indexes = list()
    new_values = list()
    df = df.dropna(subset=[column])
    for i, presplit in enumerate(df[column].astype(str)):
        values = presplit.split(sep)
        if keep and len(values) > 1:
            indexes.append(i)
            new_values.append(presplit)
        for value in values:
            indexes.append(i)
            new_values.append(value)
    new_df = df.iloc[indexes, :].copy()
    new_df[column] = new_values
    return new_df

max_stars_repo_name: szymonos/python-functions

max_issues_repo_name: szymonos/python-functions

max_forks_repo_name: szymonos/python-functions