text: # Copyright 2020 Huawei Technologies Co., Ltd
  #
  # Licensed under the Apache License, Version 2.0 (the "License");
  # you may not use this file except in compliance with the License.
  # You may obtain a copy of the License at
  #
  # http://www.apache.org/licenses/LICENSE-2.0
  #
  # Unless required by applicable law or agreed to in writing, software
  # distributed under the License is distributed on an "AS IS" BASIS,
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  # See the License for the specific language governing permissions and
  # limitations under the License.
  # ============================================================================
  """Face Recognition eval."""
  import os
  import re
  import warnings
  import argparse
  import numpy as np
  from PIL import Image
  from tqdm import tqdm
  
  import mindspore.dataset.vision.py_transforms as V
  import mindspore.dataset.transforms.py_transforms as T
  from mindspore import context, Tensor
  from mindspore.train.serialization import load_checkpoint, load_param_into_net
  
  from src.reid import SphereNet
  
  warnings.filterwarnings('ignore')
  devid = int(os.getenv('DEVICE_ID'))
  context.set_context(mode=context.GRAPH_MODE, device_target="Ascend", save_graphs=True, device_id=devid)
  
  
  def inclass_likehood(ims_info, types='cos'):
      '''Inclass likehood.'''
      obj_feas = {}
      likehoods = []
      for name, _, fea in ims_info:
          if re.split('_\\d\\d\\d\\d', name)[0] not in obj_feas:
              obj_feas[re.split('_\\d\\d\\d\\d', name)[0]] = []
          obj_feas[re.split('_\\d\\d\\d\\d', name)[0]].append(fea) # pylint: "_\d\d\d\d" -> "_\\d\\d\\d\\d"
      for _, feas in tqdm(obj_feas.items()):
          feas = np.array(feas)
          if types == 'cos':
              likehood_mat = np.dot(feas, np.transpose(feas)).tolist()
              for row in likehood_mat:
                  likehoods += row
          else:
              for fea in feas.tolist():
                  likehoods += np.sum(-(fea - feas) ** 2, axis=1).tolist()
  
      likehoods = np.array(likehoods)
      return likehoods
  
  
  def btclass_likehood(ims_info, types='cos'):
      '''Btclass likehood.'''
      likehoods = []
      count = 0
      for name1, _, fea1 in tqdm(ims_info):
          count += 1
          # pylint: "_\d\d\d\d" -> "_\\d\\d\\d\\d"
          frame_id1, _ = re.split('_\\d\\d\\d\\d', name1)[0], name1.split('_')[-1]
          fea1 = np.array(fea1)
          for name2, _, fea2 in ims_info:
              # pylint: "_\d\d\d\d" -> "_\\d\\d\\d\\d"
              frame_id2, _ = re.split('_\\d\\d\\d\\d', name2)[0], name2.split('_')[-1]
              if frame_id1 == frame_id2:
                  continue
              fea2 = np.array(fea2)
              if types == 'cos':
                  likehoods.append(np.sum(fea1 * fea2))
              else:
                  likehoods.append(np.sum(-(fea1 - fea2) ** 2))
  
      likehoods = np.array(likehoods)
      return likehoods
  
  
  def tar_at_far(inlikehoods, btlikehoods):
      test_point = [0.5, 0.3, 0.1, 0.01, 0.001, 0.0001, 0.00001]
      tar_far = []
      for point in test_point:
          thre = btlikehoods[int(btlikehoods.size * point)]
          n_ta = np.sum(inlikehoods > thre)
          tar_far.append((point, float(n_ta) / inlikehoods.size, thre))
  
      return tar_far
  
  
  def load_images(paths, batch_size=128):
      '''Load images.'''
      ll = []
      resize = V.Resize((96, 64))
      transform = T.Compose([
          V.ToTensor(),
          V.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])])
      for i, _ in enumerate(paths):
          im = Image.open(paths[i])
          im = resize(im)
          img = np.array(im)
          ts = transform(img)
          ll.append(ts[0])
          if len(ll) == batch_size:
              yield np.stack(ll, axis=0)
              ll.clear()
      if ll:
          yield np.stack(ll, axis=0)
  
  
  def main(args):
      model_path = args.pretrained
      result_file = model_path.replace('.ckpt', '.txt')
      if os.path.exists(result_file):
          os.remove(result_file)
  
      with open(result_file, 'a+') as result_fw:
          result_fw.write(model_path + '\n')
  
          network = SphereNet(num_layers=12, feature_dim=128, shape=(96, 64))
          if os.path.isfile(model_path):
              param_dict = load_checkpoint(model_path)
              param_dict_new = {}
              for key, values in param_dict.items():
                  if key.startswith('moments.'):
                      continue
                  elif key.startswith('model.'):
                      param_dict_new[key[6:]] = values
                  else:
                      param_dict_new[key] = values
              load_param_into_net(network, param_dict_new)
              print('-----------------------load model success-----------------------')
          else:
              print('-----------------------load model failed -----------------------')
  
          network.add_flags_recursive(fp16=True)
          network.set_train(False)
  
          root_path = args.eval_dir
          root_file_list = os.listdir(root_path)
          ims_info = []
          for sub_path in root_file_list:
              for im_path in os.listdir(os.path.join(root_path, sub_path)):
                  ims_info.append((im_path.split('.')[0], os.path.join(root_path, sub_path, im_path)))
  
          paths = [path for name, path in ims_info]
          names = [name for name, path in ims_info]
          print("exact feature...")
  
          l_t = []
          for batch in load_images(paths):
              batch = batch.astype(np.float32)
              batch = Tensor(batch)
              fea = network(batch)
              l_t.append(fea.asnumpy().astype(np.float16))
          feas = np.concatenate(l_t, axis=0)
          ims_info = list(zip(names, paths, feas.tolist()))
  
          print("exact inclass likehood...")
          inlikehoods = inclass_likehood(ims_info)
          inlikehoods[::-1].sort()
  
          print("exact btclass likehood...")
          btlikehoods = btclass_likehood(ims_info)
          btlikehoods[::-1].sort()
          tar_far = tar_at_far(inlikehoods, btlikehoods)
  
          for far, tar, thre in tar_far:
              print('---{}: {}@{}'.format(far, tar, thre))
  
          for far, tar, thre in tar_far:
              result_fw.write('{}: {}@{} \n'.format(far, tar, thre))
  
  
  if __name__ == "__main__":
      parser = argparse.ArgumentParser(description='reid test')
      parser.add_argument('--pretrained', type=str, default='', help='pretrained model to load')
      parser.add_argument('--eval_dir', type=str, default='', help='eval image dir, e.g. /home/test')
  
      arg = parser.parse_args()
      print(arg)
  
      main(arg)
  