def crop_face(args):
assert osp.exists(args.data_dir), "The input dir not exist"
root_folder_name = args.data_dir.split('/')[-1]
dst_folder = args.data_dir.replace(root_folder_name, root_folder_name + '_OPPOFaces')
lz.mkdir_p(dst_folder, delete=False)
mtcnn = MTCNN()
ind = 0
all_img = []
for imgfn in itertools.chain(
glob.glob(args.data_dir + '/**/*.jpg', recursive=True),
glob.glob(args.data_dir + '/**/*.JPEG', recursive=True)):
ind += 1
if ind % 10 == 0:
print(f'proc {ind}, {imgfn}')
dstimgfn = imgfn.replace(root_folder_name, root_folder_name + '_OPPOFaces')
dst_folder = osp.dirname(dstimgfn)
lz.mkdir_p(dst_folder, delete=False)
img = cvb.read_img(imgfn) # bgr
img1 = Image.fromarray(img)
face = mtcnn.align_best(img1, limit=None, min_face_size=16, imgfn=imgfn)
face = np.asarray(face) # bgr
# face = cvb.bgr2rgb(face) # rgb
cvb.write_img(face, dstimgfn)
all_img.append( dstimgfn)
logging.info(f'finish crop all {ind} imgs')
lz.msgpack_dump(all_img, dst_folder + '/' + 'all_imgs.pk')
del mtcnn
torch.cuda.empty_cache()
def extract_ms1m_info():
self = edict()
path_ms1m = lz.share_path2 + 'faces_ms1m_112x112/'
path_imgrec = lz.share_path2 + 'faces_ms1m_112x112/train.rec'
path_imgidx = path_imgrec[0:-4] + ".idx"
self.imgrec = recordio.MXIndexedRecordIO(path_imgidx, path_imgrec, 'r')
s = self.imgrec.read_idx(0)
header, _ = recordio.unpack(s)
self.header0 = (int(header.label[0]), int(header.label[1]))
# assert(header.flag==1)
self.imgidx = list(range(1, int(header.label[0])))
id2range = dict()
self.seq_identity = list(range(int(header.label[0]), int(header.label[1])))
for identity in self.seq_identity:
s = self.imgrec.read_idx(identity)
header, _ = recordio.unpack(s)
a, b = int(header.label[0]), int(header.label[1])
id2range[(identity - 3804847)] = (a, b)
count = b - a
self.seq = self.imgidx
self.seq_identity = [int(t) - 3804847 for t in self.seq_identity]
lz.msgpack_dump([self.imgidx, self.seq_identity, id2range],
path_ms1m + '/info.pk')
img = img.copy()[..., ::-1].reshape(112, 112, 3)
img = Image.fromarray(img)
mirror = transforms.functional.hflip(img)
with torch.no_grad():
fea = learner.model(conf.test_transform(img).cuda().unsqueeze(0))
fea_mirror = learner.model(conf.test_transform(mirror).cuda().unsqueeze(0))
fea = l2_norm(fea + fea_mirror).cpu().numpy().reshape(512)
return fea
def extract_fea(res):
res3 = {}
for path, img in res.items():
res3[path] = extract_fea_from_img(img)
return res3
# res, res2 = lz.msgpack_load(lz.work_path + '/yy.yy2.pk')
# res3 = extract_fea(res)
# res4 = extract_fea(res2)
# lz.msgpack_dump([res3, res4], lz.work_path + 'yy.yy2.fea.pk')
imgs = pims.ImageSequence(lz.work_path + 'face.yy2/gallery/*.png', process_func=cvb.rgb2bgr)
res = {}
for img, path in zip(imgs, imgs._filepaths):
print(path)
fea = extract_fea_from_img(img)
res[path] = fea
lz.msgpack_dump(res, lz.work_path + '/face.yy2/gallery/gallery.pk')
def crop_face(args):
for k, v in default_args.items():
setattr(args, k, v)
assert osp.exists(args.data_dir), "The input dir not exist"
root_folder_name = args.data_dir.split('/')[-1]
src_folder = args.data_dir
dst_folder = args.data_dir.replace(root_folder_name, root_folder_name + '_OPPOFaces')
lz.mkdir_p(dst_folder, delete=False)
ds = TestData(src_folder)
loader = torch.utils.data.DataLoader(ds, batch_size=args.batch_size,
num_workers=args.num_workers,
shuffle=False,
pin_memory=True,
drop_last=False
)
# 1. load pre-tained model
checkpoint_fp = 'models/phase1_wpdc_vdc.pth.tar'
arch = 'mobilenet_1'
checkpoint = torch.load(checkpoint_fp, map_location=lambda storage, loc: storage)['state_dict']
model = getattr(mobilenet_v1, arch)(num_classes=62) # 62 = 12(pose) + 40(shape) +10(expression)
model_dict = model.state_dict()
# because the model is trained by multiple gpus, prefix module should be removed
for k in checkpoint.keys():
model_dict[k.replace('module.', '')] = checkpoint[k]
model.load_state_dict(model_dict)
if args.mode == 'gpu':
cudnn.benchmark = True
model = model.cuda()
model.eval()
# 2. load dlib model for face detection and landmark used for face cropping
queue = Queue()
lock = Lock()
consumers = []
for i in range(args.num_consumers):
p = Process(target=consumer, args=(queue, lock))
p.daemon = True
consumers.append(p)
for c in consumers:
c.start()
# 3. forward
ttl_nimgs = 0
ttl_imgs = []
data_meter = lz.AverageMeter()
model_meter = lz.AverageMeter()
post_meter = lz.AverageMeter()
lz.timer.since_last_check('start crop face')
for ind, data in enumerate(loader):
data_meter.update(lz.timer.since_last_check(verbose=False))
if (data['finish'] == 1).all().item():
logging.info('finish')
break
if ind % 10 == 0:
logging.info(
f'proc batch {ind}, data time: {data_meter.avg:.2f}, model: {model_meter.avg:.2f}, post: {post_meter.avg:.2f}')
mask = data['finish'] == 0
input = data['img'][mask]
input_np = input.numpy()
roi_box = data['roi_box'][mask].numpy()
imgfn = np.asarray(data['imgfn'])[mask.numpy().astype(bool)]
dst_imgfn = [img_fp.replace(root_folder_name, root_folder_name + '_OPPOFaces') for img_fp in imgfn]
ttl_imgs.extend(dst_imgfn)
ttl_nimgs += mask.sum().item()
with torch.no_grad():
if args.mode == 'gpu':
input = input.cuda()
param = model(input)
param = param.squeeze().cpu().numpy().astype(np.float32)
model_meter.update(lz.timer.since_last_check(verbose=False))
queue.put((imgfn, param, roi_box, dst_imgfn))
# pts68 = [predict_68pts(param[i], roi_box[i]) for i in range(param.shape[0])]
# pts68_proc = [predict_68pts(param[i], [0, 0, STD_SIZE, STD_SIZE]) for i in range(param.shape[0])]
# for img_fp, pts68_, pts68_proc_, img_, dst in zip(imgfn, pts68, pts68_proc, input_np, dst_imgfn):
# ## this may need opt to async read write
# img_ori = cvb.read_img(img_fp)
# pts5 = to_landmark5(pts68_[:2, :].transpose())
# warped = preprocess(img_ori, landmark=pts5)
# # plt_imshow(warped, inp_mode='bgr'); plt.show()
# lz.mkdir_p(osp.dirname(dst), delete=False)
# cvb.write_img(warped, dst)
#
# ## this may cause black margin
# # pts5 = to_landmark5(pts68_proc_[:2, :].transpose())
# # warped = preprocess(to_img(img_), landmark=pts5)
# # # plt_imshow(warped, inp_mode='bgr'); plt.show()
# # dst = img_fp.replace(root_folder_name, root_folder_name + '_OPPOFaces')
# # cvb.write_img(warped, dst)
# if args.dump_res:
# img_ori = cvb.read_img(img_fp)
# pts_res = [pts68_]
# dst = img_fp.replace(root_folder_name, root_folder_name + '_kpts.demo')
# lz.mkdir_p(osp.dirname(dst), delete=False)
# draw_landmarks(img_ori, pts_res,
# wfp=dst,
# show_flg=args.show_flg)
post_meter.update(lz.timer.since_last_check(verbose=False))
lz.msgpack_dump(ttl_imgs, dst_folder + '/' + 'all_imgs.pk')
del model, input
torch.cuda.empty_cache()
while not queue.empty():
time.sleep(1)
def test_ijbc3(conf, learner):
if not use_mxnet:
learner.model.eval()
ds = DatasetIJBC2(flip=False)
loader = torch.utils.data.DataLoader(
ds,
batch_size=bs,
num_workers=conf.num_workers,
shuffle=False,
pin_memory=False,
)
len(loader) * bs
for ind, data in enumerate(loader):
(img, faceness_score, items, names) = data
if ind % 9 == 0:
logging.info(f'ok {ind} {len(loader)}')
if not use_mxnet:
with torch.no_grad():
img_feat = learner.model(img)
img_featf = learner.model(img.flip((3, )))
fea = (img_feat + img_featf) / 2.
fea = fea.cpu().numpy()
else:
img = img.numpy()
img_feat = learner.gets(img)
img_featf = learner.gets(img[:, :, :, ::-1].copy())
fea = (img_feat + img_featf) / 2.
# fea = fea * faceness_score.numpy().reshape(-1, 1) # todo need?
img_feats[ind * bs:min((ind + 1) * bs, ind * bs +
fea.shape[0]), :] = fea
print('last fea shape', fea.shape,
np.linalg.norm(fea, axis=-1)[-10:], img_feats.shape)
if dump_mid_res:
import h5py
f = h5py.File('/tmp/feas.ijbb.h5', 'w')
f['feas'] = img_feats
f.flush()
f.close()
templates, medias = df_tm.values[:, 1], df_tm.values[:, 2]
p1, p2, label = df_pair.values[:, 0], df_pair.values[:,
1], df_pair.values[:,
2]
unique_templates = np.unique(templates)
# cache_tfn = lz.work_path + 'ijbc.tfeas.256.pk'
cache_tfn = None
if cache_tfn and osp.exists(cache_tfn):
template_norm_feats = lz.msgpack_load(cache_tfn).copy()
clst = msgpack_load(lz.work_path + 'ijbc.tclst.pk').copy()
# clst = msgpack_load(work_path + 'ijbc.tgt.pk').copy()
lbs = np.unique(clst)
for lb in lbs:
if lb == -1: continue
mask = clst == lb
if mask.sum() <= 4: continue
nowf = template_norm_feats[mask, :]
dist = cdist(nowf, nowf)
dist[np.arange(dist.shape[0]),
np.arange(dist.shape[0])] = dist.max() * 10
wei = lz.softmax_th(-dist, dim=1, temperature=1)
refinef = np.matmul(wei, nowf)
refinef = refinef * 0.3 + nowf * 0.7
refinef = normalize(refinef, axis=1)
template_norm_feats[mask, :] = refinef
if lb % 999 == 1:
print('now refine ', lb, len(lbs), np.linalg.norm(nowf,
axis=1),
np.linalg.norm(refinef, axis=1))
else:
template_feats = np.zeros((len(unique_templates), img_feats.shape[1]))
for count_template, uqt in enumerate(unique_templates):
(ind_t, ) = np.where(templates == uqt)
face_norm_feats = img_feats[ind_t]
face_medias = medias[ind_t]
unique_medias, unique_media_counts = np.unique(face_medias,
return_counts=True)
media_norm_feats = []
for u, ct in zip(unique_medias, unique_media_counts):
(ind_m, ) = np.where(face_medias == u)
if ct == 1:
media_norm_feats += [face_norm_feats[ind_m]]
else: # image features from the same video will be aggregated into one feature
media_norm_feats += [
np.mean(face_norm_feats[ind_m], 0, keepdims=True)
]
media_norm_feats = np.array(media_norm_feats)
# media_norm_feats = media_norm_feats / np.sqrt(np.sum(media_norm_feats ** 2, -1, keepdims=True))
template_feats[count_template] = np.sum(media_norm_feats, 0)
if count_template % 2000 == 0:
print('Finish Calculating {} template features.'.format(
count_template))
template_norm_feats = template_feats / np.sqrt(
np.sum(template_feats**2, -1, keepdims=True))
if cache_tfn:
lz.msgpack_dump(template_norm_feats, cache_tfn)
template2id = np.zeros((max(unique_templates) + 1, 1), dtype=int)
for count_template, uqt in enumerate(unique_templates):
template2id[uqt] = count_template
score = np.zeros((len(p1), )) # save cosine distance between pairs
total_pairs = np.array(range(len(p1)))
batchsize = 100000 # small batchsize instead of all pairs in one batch due to the memory limiation
sublists = [
total_pairs[i:i + batchsize] for i in range(0, len(p1), batchsize)
]
total_sublists = len(sublists)
for c, s in enumerate(sublists):
feat1 = template_norm_feats[template2id[p1[s]]].reshape(-1, DIM)
feat2 = template_norm_feats[template2id[p2[s]]].reshape(-1, DIM)
# similarity_score = np.sum(feat1 * feat2, -1)
similarity_score = -(np.linalg.norm(feat1, axis=-1) + np.linalg.norm(
feat2, axis=-1) - 2 * np.sum(feat1 * feat2, -1))
score[s] = similarity_score.flatten()
if c % 10 == 0:
print('Finish {}/{} pairs.'.format(c, total_sublists))
if dump_mid_res:
msgpack_dump([label, score], f'/tmp/score.ijbb.pk')
print('score range', score.max(), score.min())
# _ = plt.hist(score)
fpr, tpr, _ = roc_curve(label, score)
# plt.figure()
# plt.semilogx(fpr, tpr, '.-')
# plt.show()
fpr = np.flipud(fpr)
tpr = np.flipud(tpr) # select largest tpr at same fpr
res = []
x_labels = [
10**-6,
10**-4,
10**-3,
]
for fpr_iter in np.arange(len(x_labels)):
_, min_index = min(
list(zip(abs(fpr - x_labels[fpr_iter]), range(len(fpr)))))
print(x_labels[fpr_iter], tpr[min_index])
res.append((x_labels[fpr_iter], tpr[min_index]))
roc_auc = auc(fpr, tpr)
print('roc aux', roc_auc)
logging.info(f'perf {res}')
if not use_mxnet:
learner.model.train()
return res
def main(args):
sys.stdout = Logger(osp.join(args.logs_dir, 'log.txt'))
sys.stderr = Logger(osp.join(args.logs_dir, 'err.txt'))
lz.init_dev(args.gpu)
print('config is {}'.format(vars(args)))
if args.seed is not None:
np.random.seed(args.seed)
torch.manual_seed(args.seed)
cudnn.benchmark = True
# Create data loaders
assert args.num_instances > 1, "num_instances should be greater than 1"
assert args.batch_size % args.num_instances == 0, \
'num_instances should divide batch_size'
test_loader = get_data(args)
# Create model
model = models.create(
args.arch,
dropout=args.dropout,
pretrained=args.pretrained,
block_name=args.block_name,
block_name2=args.block_name2,
num_features=args.num_classes,
num_classes=100,
num_deform=args.num_deform,
fusion=args.fusion,
last_conv_stride=args.last_conv_stride,
last_conv_dilation=args.last_conv_dilation,
)
print(model)
param_mb = sum(p.numel() for p in model.parameters()) / 1000000.0
print(' Total params: %.2fM' % (param_mb))
# Load from checkpoint
start_epoch = best_top1 = 0
if args.resume:
while not osp.exists(args.resume):
lz.logging.warning(' no chkpoint {} '.format(args.resume))
time.sleep(20)
if torch.cuda.is_available():
checkpoint = load_checkpoint(args.resume)
else:
checkpoint = load_checkpoint(args.resume, map_location='cpu')
# model.load_state_dict(checkpoint['state_dict'])
db_name = args.logs_dir + '/' + args.logs_dir.split('/')[-1] + '.h5'
load_state_dict(model, checkpoint['state_dict'])
with lz.Database(db_name) as db:
if 'cent' in checkpoint:
db['cent'] = to_numpy(checkpoint['cent'])
db['xent'] = to_numpy(checkpoint['state_dict']['embed2.weight'])
if args.restart:
start_epoch_ = checkpoint['epoch']
best_top1_ = checkpoint['best_top1']
print("=> Start epoch {} best top1 {:.1%}".format(
start_epoch_, best_top1_))
else:
start_epoch = checkpoint['epoch']
best_top1 = checkpoint['best_top1']
print("=> Start epoch {} best top1 {:.1%}".format(
start_epoch, best_top1))
if args.gpu is None or len(args.gpu) == 0:
model = nn.DataParallel(model)
elif len(args.gpu) == 1:
model = nn.DataParallel(model).cuda()
else:
model = nn.DataParallel(model, device_ids=range(len(args.gpu))).cuda()
# Distance metric
metric = DistanceMetric(algorithm=args.dist_metric)
# Evaluator
features, _ = extract_features(model, test_loader)
for k in features.keys():
features[k] = features[k].numpy()
lz.msgpack_dump(features, work_path + '/reid.person/fea.mp', allow_np=True)
def __init__(self, batch_size, data_shape,
path_imgrec=None,
shuffle=False, aug_list=None, mean=None,
rand_mirror=False, cutoff=0, color_jittering=0,
images_filter=0,
data_name='data', label_name='softmax_label',
# metric_learning=True,
metric_learning=False,
# todo metric learning batch formation method do not affects the scale of loss ?
**kwargs):
super(FaceImageIter, self).__init__()
self.metric_learning = metric_learning
assert path_imgrec
if path_imgrec:
logging.info('loading recordio %s...',
path_imgrec)
path_imgidx = path_imgrec[0:-4] + ".idx"
self.imgrec = recordio.MXIndexedRecordIO(path_imgidx, path_imgrec,
'r') # pylint: disable=redefined-variable-type
try:
self.imgidx, self.seq_identity, self.id2range = lz.msgpack_load(
os.path.dirname(path_imgidx) + '/info.mxnet.pk', use_list=True)
except:
s = self.imgrec.read_idx(0)
header, _ = recordio.unpack(s)
if header.flag > 0:
print('header0 label', header.label)
self.header0 = (int(header.label[0]), int(header.label[1]))
# assert(header.flag==1)
# self.imgidx = list(range(1, int(header.label[0])))
self.imgidx = []
self.id2range = {}
self.seq_identity = list(range(int(header.label[0]), int(header.label[1])))
for identity in self.seq_identity:
s = self.imgrec.read_idx(identity)
header, _ = recordio.unpack(s)
a, b = int(header.label[0]), int(header.label[1])
count = b - a
if count < images_filter:
continue
self.id2range[identity] = (a, b)
self.imgidx += range(a, b)
else:
self.imgidx = list(self.imgrec.keys)
lz.msgpack_dump([self.imgidx, self.seq_identity, self.id2range],
os.path.dirname(path_imgidx) + '/info.mxnet.pk')
print('id2range', len(self.id2range))
if shuffle:
self.seq = self.imgidx
self.oseq = self.imgidx
print(len(self.seq))
else:
self.seq = None
self.mean = mean
self.nd_mean = None
if self.mean:
self.mean = np.array(self.mean, dtype=np.float32).reshape(1, 1, 3)
self.nd_mean = mx.nd.array(self.mean).reshape((1, 1, 3))
self.check_data_shape(data_shape)
self.provide_data = [(data_name, (batch_size,) + data_shape)]
self.batch_size = batch_size
self.data_shape = data_shape
self.shuffle = shuffle
self.image_size = '%d,%d' % (data_shape[1], data_shape[2])
self.rand_mirror = rand_mirror
print('rand_mirror', rand_mirror)
self.cutoff = cutoff
self.color_jittering = color_jittering
self.CJA = mx.image.ColorJitterAug(0.125, 0.125, 0.125)
self.provide_label = [(label_name, (batch_size,))]
logging.info(f'one iter: provide {self.provide_data}, {self.provide_label}')
# print(self.provide_label[0][1])
self.cur = 0
self.nbatch = 0
self.is_init = False
self.num_instances = 4
self.num_pids_per_batch = self.batch_size // self.num_instances
self.inds_queue = []