import tensorflow as tf

# 学習済みDELGモデルを読み込みます
SAVED_MODEL_DIR = '../input/delg-saved-models/local_and_global'
DELG_MODEL = tf.saved_model.load(SAVED_MODEL_DIR)

# パラメータを設定します

## 抽出するGlobal featureの次元
NUM_EMBEDDING_DIMENSIONS = 2048
## 画像の解像度を設定します(Image Pyramids)
DELG_IMAGE_SCALES_TENSOR = tf.convert_to_tensor([0.70710677, 1.0, 1.4142135])
## DELGモデルのうち、global feature extractionに用いる部分のみ切り出します
DELG_INPUT_TENSOR_NAMES = ['input_image:0', 'input_scales:0']
GLOBAL_FEATURE_EXTRACTION_FN = DELG_MODEL.prune(DELG_INPUT_TENSOR_NAMES,
                                                ['global_descriptors:0'])


# 画像のパスから画像データをテンソルとして読み込みます
# image_pathは適宜設定してください
image_tensor = load_image_tensor(image_path)

# DELFでGlobal featureを抽出します
embedding_tensor_1 = GLOBAL_FEATURE_EXTRACTION_FN(image_tensor, DELG_IMAGE_SCALES_TENSOR)[0]

# 標準化します
embedding_tensor_2 = tf.nn.l2_normalize(
    embedding_tensor_1,
    axis=1,
    name='l2_normalization')

# 異なる解像度の結果を合体します
embedding_tensor_3 = tf.reduce_sum(
    embedding_tensor_2, axis=0, name='sum_pooling')

# さらにこれを標準化します
embedding_res = tf.nn.l2_normalize(
    embedding_tensor_3, axis=0, name='final_l2_normalization').numpy()

# パラメータを設定します
## 抽出する特徴量数の最大値
LOCAL_FEATURE_NUM_TENSOR = tf.constant(1000)

# モデルを切り出します
LOCAL_FEATURE_EXTRACTION_FN = DELG_MODEL.prune(
    DELG_INPUT_TENSOR_NAMES + ['input_max_feature_num:0', 'input_abs_thres:0'],
    ['boxes:0', 'features:0'])

# 画像をテンソルで読み込みます
image_tensor = load_image_tensor(image_path)

# DELFによるlocal featureの抽出
# 特徴量の位置と値を出力
features = LOCAL_FEATURE_EXTRACTION_FN(image_tensor, DELG_IMAGE_SCALES_TENSOR,
                                       LOCAL_FEATURE_NUM_TENSOR,
                                       DELG_SCORE_THRESHOLD_TENSOR,
                                      )

# 特徴量として抽出した画像の位置
# 出力の0列目と2列目、1列目と3列目を足して2で割る
keypoints = tf.divide(
  tf.add(
      tf.gather(features[0], [0, 1], axis=1),
      tf.gather(features[0], [2, 3], axis=1)), 2.0).numpy()


# 特徴量
# 標準化する
descriptors = tf.nn.l2_normalize(
  features[1], axis=1, name='l2_normalization').numpy()

train_keypoints = keypoints
train_descriptors = descriptors
test_keypoints = keypoints_2 #比較対象
test_descriptors = descriptors_2 #比較対象

# descriptorsを木構造に変換し、距離が近い点を取得。
# 参考：https://myenigma.hatenablog.com/entry/2020/06/14/205753#kdtree%E3%81%A8%E3%81%AF
train_descriptor_tree = spatial.cKDTree(train_descriptors)
_, matches = train_descriptor_tree.query(
  test_descriptors, distance_upper_bound=max_distance)

test_kp_count = test_keypoints.shape[0]
train_kp_count = train_keypoints.shape[0]

test_matching_keypoints = np.array([
test_keypoints[i,]
  for i in range(test_kp_count)
　#最近傍の点がmax_distanceを超える場合は最大インデックス＋１が返ってくる
  if matches[i] != train_kp_count 
])
train_matching_keypoints = np.array([
  train_keypoints[matches[i],]
  for i in range(test_kp_count)
  if matches[i] != train_kp_count
])