IJBC_Verification.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import cv2\n",
    "from glob import glob\n",
    "from matplotlib import pyplot as plt\n",
    "import matplotlib.patches as patches\n",
    "import pandas as pd\n",
    "from pathlib import Path\n",
    "import re\n",
    "from tensorflow.keras.callbacks import EarlyStopping, ModelCheckpoint, ReduceLROnPlateau, TensorBoard\n",
    "from models import ArcFaceModel\n",
    "from losses import softmax_loss\n",
    "import dataset\n",
    "import tensorflow as tf\n",
    "import os\n",
    "import logging\n",
    "import tensorflow.keras.backend as K\n",
    "from tensorflow.keras.models import Model, load_model\n",
    "from layers import ArcMarginPenaltyLogists\n",
    "from tqdm import tqdm\n",
    "from utils import l2_norm\n",
    "import logging\n",
    "\n",
    "tf.get_logger().setLevel(logging.ERROR)\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"]=\"0\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Load Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "train_size:  4658122\n",
      "steps_per_epoch:  36391\n",
      "val_size:  1164531\n",
      "validation_steps:  9097\n"
     ]
    }
   ],
   "source": [
    "### MS1M dataset\n",
    "\n",
    "batch_size = 128 # Initially 128\n",
    "input_size = 112\n",
    "embd_shape = 512\n",
    "head_type = 'CurHead' # ''ArcHead', CosHead', 'SphereHead'\n",
    "# Backbones w/ pretrained weights:\n",
    "#     MobileNet, MobileNetV2, InceptionResNetV2, InceptionV3, ResNet50, ResNet50V2, ResNet101V2, NASNetLarge, NASNetMobile, Xception\n",
    "#     But if you are trying to use NasNet, please check this issue first: https://github.com/keras-team/keras-applications/issues/78\n",
    "#         We manually download the weight file and explicitly load it in models.py file\n",
    "# Backbones w/o pretrained weights:\n",
    "#     MobileNetV3Large, MobileNetV3Small, EfficientNetLite0~6, EfficientNetB0~7\n",
    "backbone_type = 'EfficientNetB6' \n",
    "w_decay=5e-4\n",
    "num_classes = 85742 \n",
    "dataset_len = 5822653 \n",
    "base_lr = 0.01 # initially 0.01\n",
    "epochs = 20\n",
    "save_steps = 1000\n",
    "train_size = int(0.8 * dataset_len)\n",
    "print(\"train_size: \",train_size)\n",
    "steps_per_epoch = train_size // batch_size\n",
    "print(\"steps_per_epoch: \",steps_per_epoch)\n",
    "val_size = dataset_len - train_size\n",
    "print(\"val_size: \",val_size)\n",
    "validation_steps = val_size // batch_size\n",
    "print(\"validation_steps: \",validation_steps)\n",
    "steps = 1\n",
    "is_ccrop=False\n",
    "binary_img=True\n",
    "is_Adam = False\n",
    "projection_head = False  # True\n",
    "dgx = True\n",
    "\n",
    "version = \"Check\"\n",
    "\n",
    "if dgx:\n",
    "    base_dir = \"/raid/workspace/honghee/FaceRecognition/checkpoints/w_tfidentity/\"\n",
    "    if projection_head:\n",
    "        save_name = f'ms1m_{backbone_type}_{head_type}_ProjectionHead_check/{version}/*'\n",
    "    else:\n",
    "        save_name = f'ms1m_{backbone_type}_{head_type}_check/{version}/*'\n",
    "else:\n",
    "    base_dir = \"/hd/honghee/models/checkpoints/w_tfidentity/\"\n",
    "    save_name = f'ms1m_{backbone_type}_{head_type}_check/{version}/*'\n",
    "    \n",
    "# collect loss in checkpoints\n",
    "# file_list = []\n",
    "# for files in glob(base_dir+save_name):\n",
    "#     if not files.split('/')[-1].split('l_')[-1].split('.ckpt')[0] == 'checkpoint':\n",
    "#         loss = float( files.split('/')[-1].split('l_')[-1].split('.ckpt')[0] )\n",
    "#     file_list.append( loss  )\n",
    "# file_list.sort()\n",
    "\n",
    "# load_file_name = []\n",
    "# for files in glob(base_dir+save_name):\n",
    "#     if files.split('/')[-1].split('l_')[-1].split('.ckpt')[0] == 'checkpoint':\n",
    "#         pass\n",
    "#     elif file_list[0] == float( files.split('/')[-1].split('l_')[-1].split('.ckpt')[0] ) and files.split('/')[-1].split('l_')[-1].split('.ckpt')[-1]!='.index':\n",
    "#         load_file_name = files\n",
    "# best_checkpoint = load_file_name.split('.data')[0]\n",
    "# print(\"best checkpoint: \",best_checkpoint)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "best checkpoint:  /raid/workspace/honghee/FaceRecognition/checkpoints/w_tfidentity/ms1m_EfficientNetB6_CurHead_check/SGD/e_5_l_5.26140832901001.ckpt\n",
      "Model: \"arcface_model\"\n",
      "_________________________________________________________________\n",
      "Layer (type)                 Output Shape              Param #   \n",
      "=================================================================\n",
      "input_image (InputLayer)     [(None, 112, 112, 3)]     0         \n",
      "_________________________________________________________________\n",
      "efficientnetb6 (Functional)  (None, 4, 4, 2304)        40960136  \n",
      "_________________________________________________________________\n",
      "OutputLayer (Functional)     (None, 512)               18886144  \n",
      "=================================================================\n",
      "Total params: 59,846,280\n",
      "Trainable params: 59,610,584\n",
      "Non-trainable params: 235,696\n",
      "_________________________________________________________________\n"
     ]
    }
   ],
   "source": [
    "best_checkpoint = \"/raid/workspace/honghee/FaceRecognition/checkpoints/w_tfidentity/ms1m_EfficientNetB6_CurHead_check/SGD/e_5_l_5.26140832901001.ckpt\"\n",
    "print(\"best checkpoint: \",best_checkpoint)\n",
    "weight_file = best_checkpoint\n",
    "model = ArcFaceModel(size=input_size,\n",
    "                         backbone_type=backbone_type,\n",
    "                         training=False,\n",
    "                         projection_head=projection_head)\n",
    "model.load_weights(weight_file)\n",
    "model.summary()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# For IJBC"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import pickle\n",
    "\n",
    "import matplotlib\n",
    "import pandas as pd\n",
    "\n",
    "matplotlib.use('Agg')\n",
    "import matplotlib.pyplot as plt\n",
    "import timeit\n",
    "import sklearn\n",
    "import argparse\n",
    "from sklearn.metrics import roc_curve, auc\n",
    "\n",
    "from pathlib import Path\n",
    "import sys\n",
    "import warnings\n",
    "\n",
    "import cv2\n",
    "import numpy as np\n",
    "from skimage import transform as trans"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "meta_path = \"/raid/workspace/jbpark/IJB-C/IJBC/meta/\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "def read_template_media_list(path):\n",
    "    # ijb_meta = np.loadtxt(path, dtype=str)\n",
    "    ijb_meta = pd.read_csv(path, sep=' ', header=None).values\n",
    "    templates = ijb_meta[:, 1].astype(np.int)\n",
    "    medias = ijb_meta[:, 2].astype(np.int)\n",
    "    return templates, medias"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "def read_template_pair_list(path):\n",
    "    # pairs = np.loadtxt(path, dtype=str)\n",
    "    pairs = pd.read_csv(path, sep=' ', header=None).values\n",
    "    # print(pairs.shape)\n",
    "    # print(pairs[:, 0].astype(np.int))\n",
    "    t1 = pairs[:, 0].astype(np.int)\n",
    "    t2 = pairs[:, 1].astype(np.int)\n",
    "    label = pairs[:, 2].astype(np.int)\n",
    "    return t1, t2, label"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# STEP1:  Load Meta Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Time: 0.39 s. \n"
     ]
    }
   ],
   "source": [
    "# =============================================================\n",
    "# load image and template relationships for template feature embedding\n",
    "# tid --> template id,  mid --> media id\n",
    "# format:\n",
    "#           image_name tid mid\n",
    "# =============================================================\n",
    "start = timeit.default_timer()\n",
    "templates, medias = read_template_media_list(meta_path +\"ijbc_face_tid_mid.txt\")\n",
    "stop = timeit.default_timer()\n",
    "print('Time: %.2f s. ' % (stop - start))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([     1,      1,      1, ..., 187955, 187955, 187955])"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "templates"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([ 69544,   3720,      3, ..., 111105, 111105, 111105])"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "medias"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "def template_encoding(templates, medias, img_norm_feats):\n",
    "# def template_encoding(templates, medias, img_norm_feats):\n",
    "    # ==========================================================\n",
    "    # 1. face image --> l2 normalization.\n",
    "    # 2. compute media encoding.\n",
    "    # 3. compute template encoding.\n",
    "    # 4. save template features.\n",
    "    # ==========================================================\n",
    "    print('==> compute template-level feature encoding.')\n",
    "\n",
    "    uq_temp = np.unique(templates)\n",
    "    num_temp = len(uq_temp)\n",
    "    tmp_feats = np.empty((num_temp, 512))\n",
    "\n",
    "    for c, uqt in enumerate(uq_temp):\n",
    "        (ind_t,) = np.where(templates == uqt)\n",
    "#         face_norm_feats = img_norm_feats[ind_t]\n",
    "        faces_media = medias[ind_t]\n",
    "        uqm, counts = np.unique(faces_media, return_counts=True)\n",
    "        media_norm_feats = []\n",
    "\n",
    "        for u,ct in zip(uqm, counts):\n",
    "            (ind_m,) = np.where(faces_media == u)\n",
    "            if ct < 2:\n",
    "                media_norm_feats += [face_norm_feats[ind_m]]\n",
    "            else:\n",
    "                media_norm_feats += [np.sum(face_norm_feats[ind_m], 0, keepdims=True)]\n",
    "\n",
    "        media_norm_feats = np.array(media_norm_feats)\n",
    "        media_norm_feats = media_norm_feats / np.sqrt(np.sum(media_norm_feats ** 2, -1, keepdims=True))\n",
    "        template_norm_feats = np.sum(media_norm_feats, 0)\n",
    "        tmp_feats[c] = template_norm_feats\n",
    "        if c % 500 == 0:\n",
    "            print('-> finish encoding {}/{} templates.'.format(c, num_temp))\n",
    "    return tmp_feats\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Time: 2.84 s. \n"
     ]
    }
   ],
   "source": [
    "# =============================================================\n",
    "# load template pairs for template-to-template verification\n",
    "# tid : template id,  label : 1/0\n",
    "# format:\n",
    "#           tid_1 tid_2 label\n",
    "# =============================================================\n",
    "start = timeit.default_timer()\n",
    "p1, p2, label = read_template_pair_list(meta_path +\"ijbc_template_pair_label.txt\")\n",
    "stop = timeit.default_timer()\n",
    "print('Time: %.2f s. ' % (stop - start))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([     1,      1,      1, ..., 171707, 171707, 171707])"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "p1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([ 11065,  11066,  11067, ..., 183140, 182146, 185794])"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "p2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([1, 1, 1, ..., 0, 0, 0])"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "label"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# STEP2: Get Image Features"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 68,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_embed(rimg, landmark, input_size):\n",
    "    image_size = (input_size,input_size)\n",
    "\n",
    "    img = cv2.resize(rimg, dsize=(image_size[1], image_size[0]))\n",
    "    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)\n",
    "    img_flip = np.fliplr(img)\n",
    "#     img = np.transpose(img, (2, 0, 1))  # 3*112*112, RGB\n",
    "#     img_flip = np.transpose(img_flip, (2, 0, 1))\n",
    "    input_blob = np.zeros((2, image_size[1], image_size[0], 3), dtype=np.uint8)\n",
    "    input_blob[0] = ((img/255)-0.5)/0.5\n",
    "    input_blob[1] = ((img_flip/255)-0.5)/0.5\n",
    "    return input_blob"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 69,
   "metadata": {},
   "outputs": [],
   "source": [
    "# files[0]\n",
    "# name_lmk_score = files[0].strip().split(' ')\n",
    "# img_name = os.path.join(img_path, name_lmk_score[0])\n",
    "# img = cv2.imread(img_name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 70,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_image_feature(img_path, files_list, model):\n",
    "    batch_size = 2048\n",
    "#     batch_size = 128\n",
    "\n",
    "    files = files_list\n",
    "    print('files:', len(files))\n",
    "    rare_size = len(files) % batch_size\n",
    "    faceness_scores = []\n",
    "    batch = 0\n",
    "    img_feats = np.empty((len(files), 1024), dtype=np.float32)\n",
    "\n",
    "    batch_data = np.empty((2 * batch_size, input_size, input_size, 3))\n",
    "    for img_index, each_line in enumerate(files[:len(files) - rare_size]):\n",
    "        name_lmk_score = each_line.strip().split(' ')\n",
    "        img_name = os.path.join(img_path, name_lmk_score[0])\n",
    "        img = cv2.imread(img_name)\n",
    "        lmk = np.array([float(x) for x in name_lmk_score[1:-1]],\n",
    "                       dtype=np.float32)\n",
    "        lmk = lmk.reshape((5, 2))\n",
    "        input_blob = get_embed(img, lmk, input_size)\n",
    "\n",
    "        batch_data[2 * (img_index - batch * batch_size)][:] = input_blob[0]\n",
    "        batch_data[2 * (img_index - batch * batch_size) + 1][:] = input_blob[1]\n",
    "        if (img_index + 1) % batch_size == 0:\n",
    "            print('batch', batch)\n",
    "            img_feats[batch * batch_size:batch * batch_size +\n",
    "                      batch_size][:] = model.predict(batch_data).reshape([batch_size,2*512])\n",
    "#             img_feats[batch * batch_size:batch * batch_size +\n",
    "#                       batch_size][:] = tf.reshape(model(batch_data), [batch_size,2*512])\n",
    "            batch += 1\n",
    "        faceness_scores.append(name_lmk_score[-1])\n",
    "\n",
    "    batch_data = np.empty((2 * rare_size, input_size, input_size, 3))\n",
    "    \n",
    "    # batch size로 안나눠지는 남는 것들 처리\n",
    "    for img_index, each_line in enumerate(files[len(files) - rare_size:]):\n",
    "        name_lmk_score = each_line.strip().split(' ')\n",
    "        img_name = os.path.join(img_path, name_lmk_score[0])\n",
    "        img = cv2.imread(img_name)\n",
    "        lmk = np.array([float(x) for x in name_lmk_score[1:-1]],\n",
    "                       dtype=np.float32)\n",
    "        lmk = lmk.reshape((5, 2))\n",
    "        input_blob = get_embed(img, lmk, input_size)\n",
    "        batch_data[2 * img_index][:] = input_blob[0]\n",
    "        batch_data[2 * img_index + 1][:] = input_blob[1]\n",
    "        if (img_index + 1) % rare_size == 0:\n",
    "            print('batch', batch)\n",
    "            img_feats[len(files) - rare_size:][:] = model.predict(batch_data).reshape([rare_size,2*512])\n",
    "#             img_feats[len(files) - rare_size:][:] = tf.reshape(model(batch_data), [rare_size,2*512])\n",
    "            batch += 1\n",
    "        faceness_scores.append(name_lmk_score[-1])\n",
    "    faceness_scores = np.array(faceness_scores).astype(np.float32)\n",
    "    # img_feats = np.ones( (len(files), 1024), dtype=np.float32) * 0.01\n",
    "    # faceness_scores = np.ones( (len(files), ), dtype=np.float32 )\n",
    "    return img_feats, faceness_scores\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 71,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "/raid/workspace/jbpark/IJB-C/veri_crops\n"
     ]
    }
   ],
   "source": [
    "# =============================================================\n",
    "# load image features\n",
    "# format:\n",
    "#           img_feats: [image_num x feats_dim] (227630, 512)\n",
    "# =============================================================\n",
    "start = timeit.default_timer()\n",
    "image_path = \"/raid/workspace/jbpark/IJB-C\"\n",
    "img_path = '%s/veri_crops' % image_path\n",
    "print(img_path)\n",
    "img_list_path = meta_path +\"ijbc_name_5pts_score.txt\"\n",
    "img_list = open(img_list_path)\n",
    "files = img_list.readlines()\n",
    "# files_list = divideIntoNstrand(files, rank_size)\n",
    "files_list = files"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 72,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "files: 469375\n",
      "batch 0\n",
      "batch 1\n",
      "batch 2\n",
      "batch 3\n",
      "batch 4\n",
      "batch 5\n",
      "batch 6\n",
      "batch 7\n",
      "batch 8\n",
      "batch 9\n",
      "batch 10\n",
      "batch 11\n",
      "batch 12\n",
      "batch 13\n",
      "batch 14\n",
      "batch 15\n",
      "batch 16\n",
      "batch 17\n",
      "batch 18\n",
      "batch 19\n",
      "batch 20\n",
      "batch 21\n",
      "batch 22\n",
      "batch 23\n",
      "batch 24\n",
      "batch 25\n",
      "batch 26\n",
      "batch 27\n",
      "batch 28\n",
      "batch 29\n",
      "batch 30\n",
      "batch 31\n",
      "batch 32\n",
      "batch 33\n",
      "batch 34\n",
      "batch 35\n",
      "batch 36\n",
      "batch 37\n",
      "batch 38\n",
      "batch 39\n",
      "batch 40\n",
      "batch 41\n",
      "batch 42\n",
      "batch 43\n",
      "batch 44\n",
      "batch 45\n",
      "batch 46\n",
      "batch 47\n",
      "batch 48\n",
      "batch 49\n",
      "batch 50\n",
      "batch 51\n",
      "batch 52\n",
      "batch 53\n",
      "batch 54\n",
      "batch 55\n",
      "batch 56\n",
      "batch 57\n",
      "batch 58\n",
      "batch 59\n",
      "batch 60\n",
      "batch 61\n",
      "batch 62\n",
      "batch 63\n",
      "batch 64\n",
      "batch 65\n",
      "batch 66\n",
      "batch 67\n",
      "batch 68\n",
      "batch 69\n",
      "batch 70\n",
      "batch 71\n",
      "batch 72\n",
      "batch 73\n",
      "batch 74\n",
      "batch 75\n",
      "batch 76\n",
      "batch 77\n",
      "batch 78\n",
      "batch 79\n",
      "batch 80\n",
      "batch 81\n",
      "batch 82\n",
      "batch 83\n",
      "batch 84\n",
      "batch 85\n",
      "batch 86\n",
      "batch 87\n",
      "batch 88\n",
      "batch 89\n",
      "batch 90\n",
      "batch 91\n",
      "batch 92\n",
      "batch 93\n",
      "batch 94\n",
      "batch 95\n",
      "batch 96\n",
      "batch 97\n",
      "batch 98\n",
      "batch 99\n",
      "batch 100\n",
      "batch 101\n",
      "batch 102\n",
      "batch 103\n",
      "batch 104\n",
      "batch 105\n",
      "batch 106\n",
      "batch 107\n",
      "batch 108\n",
      "batch 109\n",
      "batch 110\n",
      "batch 111\n",
      "batch 112\n",
      "batch 113\n",
      "batch 114\n",
      "batch 115\n",
      "batch 116\n",
      "batch 117\n",
      "batch 118\n",
      "batch 119\n",
      "batch 120\n",
      "batch 121\n",
      "batch 122\n",
      "batch 123\n",
      "batch 124\n",
      "batch 125\n",
      "batch 126\n",
      "batch 127\n",
      "batch 128\n",
      "batch 129\n",
      "batch 130\n",
      "batch 131\n",
      "batch 132\n",
      "batch 133\n",
      "batch 134\n",
      "batch 135\n",
      "batch 136\n",
      "batch 137\n",
      "batch 138\n",
      "batch 139\n",
      "batch 140\n",
      "batch 141\n",
      "batch 142\n",
      "batch 143\n",
      "batch 144\n",
      "batch 145\n",
      "batch 146\n",
      "batch 147\n",
      "batch 148\n",
      "batch 149\n",
      "batch 150\n",
      "batch 151\n",
      "batch 152\n",
      "batch 153\n",
      "batch 154\n",
      "batch 155\n",
      "batch 156\n",
      "batch 157\n",
      "batch 158\n",
      "batch 159\n",
      "batch 160\n",
      "batch 161\n",
      "batch 162\n",
      "batch 163\n",
      "batch 164\n",
      "batch 165\n",
      "batch 166\n",
      "batch 167\n",
      "batch 168\n",
      "batch 169\n",
      "batch 170\n",
      "batch 171\n",
      "batch 172\n",
      "batch 173\n",
      "batch 174\n",
      "batch 175\n",
      "batch 176\n",
      "batch 177\n",
      "batch 178\n",
      "batch 179\n",
      "batch 180\n",
      "batch 181\n",
      "batch 182\n",
      "batch 183\n",
      "batch 184\n",
      "batch 185\n",
      "batch 186\n",
      "batch 187\n",
      "batch 188\n",
      "batch 189\n",
      "batch 190\n",
      "batch 191\n",
      "batch 192\n",
      "batch 193\n",
      "batch 194\n",
      "batch 195\n",
      "batch 196\n",
      "batch 197\n",
      "batch 198\n",
      "batch 199\n",
      "batch 200\n",
      "batch 201\n",
      "batch 202\n",
      "batch 203\n",
      "batch 204\n",
      "batch 205\n",
      "batch 206\n",
      "batch 207\n",
      "batch 208\n",
      "batch 209\n",
      "batch 210\n",
      "batch 211\n",
      "batch 212\n",
      "batch 213\n",
      "batch 214\n",
      "batch 215\n",
      "batch 216\n",
      "batch 217\n",
      "batch 218\n",
      "batch 219\n",
      "batch 220\n",
      "batch 221\n",
      "batch 222\n",
      "batch 223\n",
      "batch 224\n",
      "batch 225\n",
      "batch 226\n",
      "batch 227\n",
      "batch 228\n",
      "batch 229\n",
      "Time: 2140.14 s. \n",
      "Feature Shape: (469375 , 1024) .\n"
     ]
    }
   ],
   "source": [
    "# img_feats\n",
    "# for i in range(rank_size):\n",
    "img_feats, faceness_scores = get_image_feature(img_path, files_list, model)\n",
    "stop = timeit.default_timer()\n",
    "print('Time: %.2f s. ' % (stop - start))\n",
    "print('Feature Shape: ({} , {}) .'.format(img_feats.shape[0],\n",
    "                                          img_feats.shape[1]))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# STEP3: Get Template Feature "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 73,
   "metadata": {},
   "outputs": [],
   "source": [
    "def image2template_feature(img_feats=None, templates=None, medias=None):\n",
    "    # ==========================================================\n",
    "    # 1. face image feature l2 normalization. img_feats:[number_image x feats_dim]\n",
    "    # 2. compute media feature.\n",
    "    # 3. compute template feature.\n",
    "    # ==========================================================\n",
    "    unique_templates = np.unique(templates)\n",
    "    template_feats = np.zeros((len(unique_templates), img_feats.shape[1]))\n",
    "    for count_template, uqt in enumerate(unique_templates):\n",
    "\n",
    "        (ind_t,) = np.where(templates == uqt)\n",
    "        face_norm_feats = img_feats[ind_t]\n",
    "        face_medias = medias[ind_t]\n",
    "        unique_medias, unique_media_counts = np.unique(face_medias,\n",
    "                                                       return_counts=True)\n",
    "        media_norm_feats = []\n",
    "        for u, ct in zip(unique_medias, unique_media_counts):\n",
    "            (ind_m,) = np.where(face_medias == u)\n",
    "            if ct == 1:\n",
    "                media_norm_feats += [face_norm_feats[ind_m]]\n",
    "            else:  # image features from the same video will be aggregated into one feature\n",
    "                media_norm_feats += [\n",
    "                    np.mean(face_norm_feats[ind_m], axis=0, keepdims=True)\n",
    "                ]\n",
    "        media_norm_feats = np.array(media_norm_feats)\n",
    "        # media_norm_feats = media_norm_feats / np.sqrt(np.sum(media_norm_feats ** 2, -1, keepdims=True))\n",
    "        template_feats[count_template] = np.sum(media_norm_feats, axis=0)\n",
    "        if count_template % 2000 == 0:\n",
    "            print('Finish Calculating {} template features.'.format(\n",
    "                count_template))\n",
    "    # template_norm_feats = template_feats / np.sqrt(np.sum(template_feats ** 2, -1, keepdims=True))\n",
    "    template_norm_feats = sklearn.preprocessing.normalize(template_feats)\n",
    "    # print(template_norm_feats.shape)\n",
    "    return template_norm_feats, unique_templates"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 74,
   "metadata": {},
   "outputs": [],
   "source": [
    "# =============================================================\n",
    "# compute template features from image features.\n",
    "# =============================================================\n",
    "start = timeit.default_timer()\n",
    "# ==========================================================\n",
    "# Norm feature before aggregation into template feature?\n",
    "# Feature norm from embedding network and faceness score are able to decrease weights for noise samples (not face).\n",
    "# ==========================================================\n",
    "# 1. FaceScore （Feature Norm）\n",
    "# 2. FaceScore （Detector）\n",
    "use_flip_test = False\n",
    "use_norm_score = True\n",
    "use_detector_score = False\n",
    "\n",
    "# if use_flip_test:\n",
    "#     # concat --- F1\n",
    "#     # img_input_feats = img_feats\n",
    "#     # add --- F2\n",
    "#     img_input_feats = img_feats[:, 0:img_feats.shape[1] //\n",
    "#                                      2] + img_feats[:, img_feats.shape[1] // 2:]\n",
    "# else:\n",
    "#     img_input_feats = img_feats[:, 0:img_feats.shape[1] // 2]\n",
    "img_input_feats = img_feats\n",
    "    \n",
    "    \n",
    "if use_norm_score:\n",
    "    img_input_feats = img_input_feats\n",
    "else:\n",
    "    # normalise features to remove norm information\n",
    "    img_input_feats = img_input_feats / np.sqrt(\n",
    "        np.sum(img_input_feats ** 2, -1, keepdims=True))\n",
    "\n",
    "if use_detector_score:\n",
    "    print(img_input_feats.shape, faceness_scores.shape)\n",
    "    img_input_feats = img_input_feats * faceness_scores[:, np.newaxis]\n",
    "else:\n",
    "    img_input_feats = img_input_feats"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 75,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Finish Calculating 0 template features.\n",
      "Finish Calculating 2000 template features.\n",
      "Finish Calculating 4000 template features.\n",
      "Finish Calculating 6000 template features.\n",
      "Finish Calculating 8000 template features.\n",
      "Finish Calculating 10000 template features.\n",
      "Finish Calculating 12000 template features.\n",
      "Finish Calculating 14000 template features.\n",
      "Finish Calculating 16000 template features.\n",
      "Finish Calculating 18000 template features.\n",
      "Finish Calculating 20000 template features.\n",
      "Finish Calculating 22000 template features.\n",
      "Time: 9.98 s. \n"
     ]
    }
   ],
   "source": [
    "template_norm_feats, unique_templates = image2template_feature(\n",
    "    img_input_feats, templates, medias)\n",
    "stop = timeit.default_timer()\n",
    "print('Time: %.2f s. ' % (stop - start))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 76,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(23124, 1024)"
      ]
     },
     "execution_count": 76,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "template_norm_feats.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 77,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(23124,)"
      ]
     },
     "execution_count": 77,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "unique_templates.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# STEP4: Get Template Similarity Scores"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 78,
   "metadata": {},
   "outputs": [],
   "source": [
    "def verification(template_norm_feats=None,\n",
    "                 unique_templates=None,\n",
    "                 p1=None,\n",
    "                 p2=None):\n",
    "    # ==========================================================\n",
    "    #         Compute set-to-set Similarity Score.\n",
    "    # ==========================================================\n",
    "    template2id = np.zeros((max(unique_templates) + 1, 1), dtype=int)\n",
    "    for count_template, uqt in enumerate(unique_templates):\n",
    "        template2id[uqt] = count_template\n",
    "\n",
    "    score = np.zeros((len(p1),))  # save cosine distance between pairs\n",
    "\n",
    "    total_pairs = np.array(range(len(p1)))\n",
    "    batchsize = 100000  # small batchsize instead of all pairs in one batch due to the memory limiation\n",
    "    sublists = [\n",
    "        total_pairs[i:i + batchsize] for i in range(0, len(p1), batchsize)\n",
    "    ]\n",
    "    total_sublists = len(sublists)\n",
    "    for c, s in enumerate(sublists):\n",
    "        feat1 = template_norm_feats[template2id[p1[s]]]\n",
    "        feat2 = template_norm_feats[template2id[p2[s]]]\n",
    "        similarity_score = np.sum(feat1 * feat2, -1)\n",
    "        score[s] = similarity_score.flatten()\n",
    "        if c % 10 == 0:\n",
    "            print('Finish {}/{} pairs.'.format(c, total_sublists))\n",
    "    return score"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 79,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Finish 0/157 pairs.\n",
      "Finish 10/157 pairs.\n",
      "Finish 20/157 pairs.\n",
      "Finish 30/157 pairs.\n",
      "Finish 40/157 pairs.\n",
      "Finish 50/157 pairs.\n",
      "Finish 60/157 pairs.\n",
      "Finish 70/157 pairs.\n",
      "Finish 80/157 pairs.\n",
      "Finish 90/157 pairs.\n",
      "Finish 100/157 pairs.\n",
      "Finish 110/157 pairs.\n",
      "Finish 120/157 pairs.\n",
      "Finish 130/157 pairs.\n",
      "Finish 140/157 pairs.\n",
      "Finish 150/157 pairs.\n",
      "Time: 276.64 s. \n"
     ]
    }
   ],
   "source": [
    "# =============================================================\n",
    "# compute verification scores between template pairs.\n",
    "# =============================================================\n",
    "start = timeit.default_timer()\n",
    "score = verification(template_norm_feats, unique_templates, p1, p2)\n",
    "stop = timeit.default_timer()\n",
    "print('Time: %.2f s. ' % (stop - start))\n",
    "\n",
    "# In[ ]:\n",
    "save_path = \"result/\"\n",
    "# save_path = result_dir + '/%s_result' % target\n",
    "\n",
    "if not os.path.exists(save_path):\n",
    "    os.makedirs(save_path)\n",
    "target = f'{head_type}_{backbone_type}'\n",
    "score_save_file = os.path.join(save_path, \"%s.npy\" % target.lower())\n",
    "np.save(score_save_file, score)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 80,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(15658489,)"
      ]
     },
     "execution_count": 80,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "score.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 81,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([0.79342157, 0.8875268 , 0.68414404, ..., 0.46842706, 0.38269398,\n",
       "       0.44070727])"
      ]
     },
     "execution_count": 81,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "score"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# STEP5: Compute ROC"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 82,
   "metadata": {},
   "outputs": [],
   "source": [
    "def compute_ROC(labels, scores, roc_path):\n",
    "    print('==> compute ROC.')\n",
    "    import sklearn.metrics as skm\n",
    "    from scipy import interpolate\n",
    "    fpr, tpr, thresholds = skm.roc_curve(labels, scores)\n",
    "    roc_auc = skm.auc(fpr, tpr)\n",
    "    print(\"AUROC: \",roc_auc)\n",
    "    fpr_levels = [1e-5, 1e-4, 1e-3, 1e-2, 1e-1]\n",
    "    f_interp = interpolate.interp1d(fpr, tpr)\n",
    "    tpr_at_fpr = [f_interp(x) for x in fpr_levels]\n",
    "    roc_txt = f'ROC_{head_type}_{backbone_type}.txt'.lower()\n",
    "    file = open(f'{save_path}{roc_txt}', 'w')\n",
    "    for (far, tar) in zip(fpr_levels, tpr_at_fpr):\n",
    "        print('TAR @ FAR = {} : {}'.format(far, tar))\n",
    "        file.write('TAR @ FAR = {}: {}\\n'.format(far, tar))\n",
    "    file.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 83,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "==> compute ROC.\n",
      "AUROC:  0.5866860392641152\n",
      "TAR @ FAR = 1e-05 : 2.101858696957359e-05\n",
      "TAR @ FAR = 0.0001 : 0.0002890131965186882\n",
      "TAR @ FAR = 0.001 : 0.002965690034258833\n",
      "TAR @ FAR = 0.01 : 0.019174720049087282\n",
      "TAR @ FAR = 0.1 : 0.1495938503192243\n"
     ]
    }
   ],
   "source": [
    "compute_ROC(label, score, save_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 84,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+-----------------------------------------------+-------+-------+--------+-------+------+-------+\n",
      "|                    Methods                    | 1e-06 | 1e-05 | 0.0001 | 0.001 | 0.01 |  0.1  |\n",
      "+-----------------------------------------------+-------+-------+--------+-------+------+-------+\n",
      "| curhead_efficientnetb6-CurHead_EfficientNetB6 |  0.00 |  0.00 |  0.00  |  0.30 | 1.92 | 14.96 |\n",
      "+-----------------------------------------------+-------+-------+--------+-------+------+-------+\n"
     ]
    }
   ],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "# import timeit\n",
    "import sklearn\n",
    "# import argparse\n",
    "from sklearn.metrics import roc_curve, auc\n",
    "\n",
    "from menpo.visualize.viewmatplotlib import sample_colours_from_colourmap\n",
    "from prettytable import PrettyTable\n",
    "from pathlib import Path\n",
    "import sys\n",
    "import warnings\n",
    "\n",
    "files = [score_save_file]\n",
    "methods = []\n",
    "scores = []\n",
    "for file in files:\n",
    "    methods.append(Path(file).stem)\n",
    "    scores.append(np.load(file))\n",
    "\n",
    "methods = np.array(methods)\n",
    "scores = dict(zip(methods, scores))\n",
    "colours = dict(\n",
    "    zip(methods, sample_colours_from_colourmap(methods.shape[0], 'Set2')))\n",
    "x_labels = [10 ** -6, 10 ** -5, 10 ** -4, 10 ** -3, 10 ** -2, 10 ** -1]\n",
    "tpr_fpr_table = PrettyTable(['Methods'] + [str(x) for x in x_labels])\n",
    "fig = plt.figure()\n",
    "for method in methods:\n",
    "    fpr, tpr, _ = roc_curve(label, scores[method])\n",
    "    roc_auc = auc(fpr, tpr)\n",
    "    fpr = np.flipud(fpr)\n",
    "    tpr = np.flipud(tpr)  # select largest tpr at same fpr\n",
    "    plt.plot(fpr,\n",
    "             tpr,\n",
    "             color=colours[method],\n",
    "             lw=1,\n",
    "             label=('[%s (AUC = %0.4f %%)]' %\n",
    "                    (method.split('-')[-1], roc_auc * 100)))\n",
    "    tpr_fpr_row = []\n",
    "    tpr_fpr_row.append(\"%s-%s\" % (method, target))\n",
    "    for fpr_iter in np.arange(len(x_labels)):\n",
    "        _, min_index = min(\n",
    "            list(zip(abs(fpr - x_labels[fpr_iter]), range(len(fpr)))))\n",
    "        tpr_fpr_row.append('%.2f' % (tpr[min_index] * 100))\n",
    "    tpr_fpr_table.add_row(tpr_fpr_row)\n",
    "plt.xlim([10 ** -6, 0.1])\n",
    "plt.ylim([0.3, 1.0])\n",
    "plt.grid(linestyle='--', linewidth=1)\n",
    "plt.xticks(x_labels)\n",
    "plt.yticks(np.linspace(0.3, 1.0, 8, endpoint=True))\n",
    "plt.xscale('log')\n",
    "plt.xlabel('False Positive Rate')\n",
    "plt.ylabel('True Positive Rate')\n",
    "plt.title('ROC on IJB')\n",
    "plt.legend(loc=\"lower right\")\n",
    "# fig.savefig(os.path.join(save_path, '%s.pdf' % target.lower()))\n",
    "print(tpr_fpr_table)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}