mirror of
https://github.com/deepinsight/insightface.git
synced 2026-05-19 15:41:33 +00:00
71 lines
2.2 KiB
Python
71 lines
2.2 KiB
Python
from __future__ import absolute_import
|
|
from __future__ import division
|
|
from __future__ import print_function
|
|
import os
|
|
import sys
|
|
import mxnet as mx
|
|
from mxnet import ndarray as nd
|
|
import random
|
|
import argparse
|
|
import cv2
|
|
import time
|
|
import sklearn
|
|
from sklearn.decomposition import PCA
|
|
from easydict import EasyDict as edict
|
|
from sklearn.cluster import DBSCAN
|
|
import numpy as np
|
|
|
|
sys.path.append(os.path.join(os.path.dirname(__file__),'..', 'common'))
|
|
import face_image
|
|
|
|
def main(args):
|
|
include_datasets = args.include.split(',')
|
|
rec_list = []
|
|
for ds in include_datasets:
|
|
path_imgrec = os.path.join(ds, 'train.rec')
|
|
path_imgidx = os.path.join(ds, 'train.idx')
|
|
imgrec = mx.recordio.MXIndexedRecordIO(path_imgidx, path_imgrec, 'r') # pylint: disable=redefined-variable-type
|
|
rec_list.append(imgrec)
|
|
if not os.path.exists(args.output):
|
|
os.makedirs(args.output)
|
|
writer = mx.recordio.MXIndexedRecordIO(os.path.join(args.output, 'train.idx'), os.path.join(args.output, 'train.rec'), 'w')
|
|
for ds_id in xrange(len(rec_list)):
|
|
id_list = []
|
|
imgrec = rec_list[ds_id]
|
|
s = imgrec.read_idx(0)
|
|
writer.write_idx(0, s)
|
|
header, _ = mx.recordio.unpack(s)
|
|
assert header.flag>0
|
|
print('header0 label', header.label)
|
|
header0 = (int(header.label[0]), int(header.label[1]))
|
|
seq_identity = range(int(header.label[0]), int(header.label[1]))
|
|
pp=0
|
|
nlabel = -1
|
|
for identity in seq_identity:
|
|
pp+=1
|
|
if pp%10==0:
|
|
print('processing id', pp)
|
|
s = imgrec.read_idx(identity)
|
|
writer.write_idx(identity, s)
|
|
header, _ = mx.recordio.unpack(s)
|
|
nlabel+=1
|
|
for _idx in xrange(int(header.label[0]), int(header.label[1])):
|
|
s = imgrec.read_idx(_idx)
|
|
_header, _content = mx.recordio.unpack(s)
|
|
nheader = mx.recordio.IRHeader(0, nlabel, _idx, 0)
|
|
s = mx.recordio.pack(nheader, _content)
|
|
writer.write_idx(_idx, s)
|
|
|
|
print('max label', nlabel)
|
|
|
|
|
|
|
|
if __name__ == '__main__':
|
|
parser = argparse.ArgumentParser(description='do dataset merge')
|
|
# general
|
|
parser.add_argument('--include', default='', type=str, help='')
|
|
parser.add_argument('--output', default='', type=str, help='')
|
|
args = parser.parse_args()
|
|
main(args)
|
|
|