insightface/recognition/SubCenter-ArcFace/drop.py

import os
import shutil
import datetime
import sys
from mxnet import ndarray as nd
import mxnet as mx
import random
import argparse
import numbers
import cv2
import time
import pickle
import sklearn
import sklearn.preprocessing
from easydict import EasyDict as edict
import numpy as np
sys.path.append(os.path.join(os.path.dirname(__file__), '..', 'common'))
from rec_builder import *

def get_embedding(args, imgrec, a, b, image_size, model):
  ocontents = []
  for idx in range(a, b):
    s = imgrec.read_idx(idx)
    ocontents.append(s)
  embeddings = None
  #print(len(ocontents))
  ba = 0
  rlabel = -1
  imgs = []
  contents = []
  while True:
    bb = min(ba+args.batch_size, len(ocontents))
    if ba>=bb:
      break
    _batch_size = bb-ba
    #_batch_size2 = max(_batch_size, args.ctx_num)
    _batch_size2 = _batch_size
    if _batch_size%args.ctx_num!=0:
        _batch_size2 = ((_batch_size//args.ctx_num)+1) * args.ctx_num
    data = np.zeros( (_batch_size2,3, image_size[0], image_size[1]) )
    count = bb-ba
    ii=0
    for i in range(ba, bb):
      header, img = mx.recordio.unpack(ocontents[i])
      contents.append(img)
      label = header.label
      if not isinstance(label, numbers.Number):
        label = label[0]
      if rlabel<0:
        rlabel = int(label)
      
      img = mx.image.imdecode(img)
      rgb = img.asnumpy()
      bgr = rgb[:,:,::-1]
      imgs.append(bgr)
      img = rgb.transpose( (2,0,1) )
      data[ii] = img
      ii+=1
    while ii<_batch_size2:
      data[ii] = data[0]
      ii+=1
    nddata = nd.array(data)
    db = mx.io.DataBatch(data=(nddata,))
    model.forward(db, is_train=False)
    net_out = model.get_outputs()
    net_out = net_out[0].asnumpy()
    if embeddings is None:
      embeddings = np.zeros( (len(ocontents), net_out.shape[1]))
    embeddings[ba:bb,:] = net_out[0:_batch_size,:]
    ba = bb
  embeddings = sklearn.preprocessing.normalize(embeddings)
  return embeddings, rlabel, contents

def main(args):
  print(args)
  image_size = (112,112)
  print('image_size', image_size)
  vec = args.model.split(',')
  prefix = vec[0]
  epoch = int(vec[1])
  print('loading',prefix, epoch)
  ctx = []
  cvd = os.environ['CUDA_VISIBLE_DEVICES'].strip()
  if len(cvd)>0:
    for i in range(len(cvd.split(','))):
      ctx.append(mx.gpu(i))
  if len(ctx)==0:
    ctx = [mx.cpu()]
    print('use cpu')
  else:
    print('gpu num:', len(ctx))
  args.ctx_num = len(ctx)
  args.batch_size *= args.ctx_num
  sym, arg_params, aux_params = mx.model.load_checkpoint(prefix, epoch)
  W = None
  i = 0
  while True:
      key = 'fc7_%d_weight'%i
      i+=1
      if key not in arg_params:
          break
      _W = arg_params[key].asnumpy()
      #_W = _W.reshape( (-1, 10, 512) )
      if W is None:
          W = _W
      else:
          W = np.concatenate( (W, _W), axis=0 )
  K = args.k
  W = sklearn.preprocessing.normalize(W)
  W = W.reshape( (-1, K, 512) )
  all_layers = sym.get_internals()
  sym = all_layers['fc1_output']
  model = mx.mod.Module(symbol=sym, context=ctx, label_names = None)
  model.bind(data_shapes=[('data', (args.ctx_num, 3, image_size[0], image_size[1]))])
  model.set_params(arg_params, aux_params)
  print('W:',W.shape)
  path_imgrec = os.path.join(args.data, 'train.rec')
  path_imgidx = os.path.join(args.data, 'train.idx')
  imgrec = mx.recordio.MXIndexedRecordIO(path_imgidx, path_imgrec, 'r')  # pylint: disable=redefined-variable-type
  id_list = []
  s = imgrec.read_idx(0)
  header, _ = mx.recordio.unpack(s)
  assert header.flag>0
  print('header0 label', header.label)
  header0 = (int(header.label[0]), int(header.label[1]))
  #assert(header.flag==1)
  imgidx = range(1, int(header.label[0]))
  id2range = {}
  a, b = int(header.label[0]), int(header.label[1])
  seq_identity = range(a,b)
  print(len(seq_identity))
  image_count = 0
  pp=0
  for wid, identity in enumerate(seq_identity):
    pp+=1
    s = imgrec.read_idx(identity)
    header, _ = mx.recordio.unpack(s)
    contents = []
    a,b = int(header.label[0]), int(header.label[1])
    _count = b-a
    id_list.append( (wid, a, b, _count) )
    image_count += _count
  pp = 0
  if not os.path.exists(args.output):
    os.makedirs(args.output)
  ret = np.zeros( (image_count, K+1), dtype=np.float32 )
  output_dir = args.output
  builder = SeqRecBuilder(output_dir)
  print(ret.shape)
  imid = 0
  da = datetime.datetime.now()
  label = 0
  num_images = 0
  cos_thresh = np.cos(np.pi*args.threshold / 180.0)
  for id_item in id_list:
    wid = id_item[0]
    pp+=1
    if pp%40==0:
      db = datetime.datetime.now()
      print('processing id', pp, (db-da).total_seconds())
      da = db
    x, _, contents = get_embedding(args, imgrec, id_item[1], id_item[2], image_size, model)
    subcenters = W[wid]
    K_stat = np.zeros( (K, ), dtype=np.int)
    for i in range(x.shape[0]):
        _x = x[i]
        sim = np.dot(subcenters, _x) # len(sim)==K
        mc = np.argmax(sim)
        K_stat[mc] += 1
    dominant_index = np.argmax(K_stat)
    dominant_center = subcenters[dominant_index]
    sim = np.dot(x, dominant_center) 
    idx = np.where(sim>cos_thresh)[0]
    num_drop = x.shape[0] - len(idx)
    if len(idx)==0:
        continue
    #print("labelid %d dropped %d, from %d to %d"% (wid, num_drop, x.shape[0], len(idx)))
    num_images += len(idx)
    for _idx in idx:
        c = contents[_idx]
        builder.add(label, c, is_image=False)
    label+=1
  builder.close()

  print('total:', num_images)
  




if __name__ == '__main__':
  parser = argparse.ArgumentParser(description='')
  # general
  parser.add_argument('--data', default='/bigdata/faces_ms1m_full', type=str, help='')
  parser.add_argument('--output', default='/bigdata/ms1m_full_k3drop075', type=str, help='')
  parser.add_argument('--model', default='../Evaluation/IJB/pretrained_models/r50-arcfacesc-msf-k3z/model,2', help='path to load model.')
  parser.add_argument('--batch-size', default=16, type=int, help='')
  parser.add_argument('--threshold', default=75, type=float, help='')
  parser.add_argument('--k', default=3, type=int, help='')
  args = parser.parse_args()
  main(args)