mirror of
https://github.com/deepinsight/insightface.git
synced 2025-12-30 08:02:27 +00:00
137 lines
4.4 KiB
Python
137 lines
4.4 KiB
Python
from __future__ import absolute_import
|
|
from __future__ import division
|
|
from __future__ import print_function
|
|
#import mxnet as mx
|
|
#from mxnet import ndarray as nd
|
|
import argparse
|
|
import cv2
|
|
import pickle
|
|
import numpy as np
|
|
import sys
|
|
import os
|
|
sys.path.append(os.path.join(os.path.dirname(__file__), '..', '..', 'common'))
|
|
sys.path.append(
|
|
os.path.join(os.path.dirname(__file__), '..', '..', 'RetinaFace'))
|
|
import face_align
|
|
from retinaface import RetinaFace
|
|
|
|
|
|
def to_rgb(img):
|
|
w, h = img.shape
|
|
ret = np.empty((w, h, 3), dtype=np.uint8)
|
|
ret[:, :, 0] = ret[:, :, 1] = ret[:, :, 2] = img
|
|
return ret
|
|
|
|
|
|
def IOU(Reframe, GTframe):
|
|
x1 = Reframe[0]
|
|
y1 = Reframe[1]
|
|
width1 = Reframe[2] - Reframe[0]
|
|
height1 = Reframe[3] - Reframe[1]
|
|
|
|
x2 = GTframe[0]
|
|
y2 = GTframe[1]
|
|
width2 = GTframe[2] - GTframe[0]
|
|
height2 = GTframe[3] - GTframe[1]
|
|
|
|
endx = max(x1 + width1, x2 + width2)
|
|
startx = min(x1, x2)
|
|
width = width1 + width2 - (endx - startx)
|
|
|
|
endy = max(y1 + height1, y2 + height2)
|
|
starty = min(y1, y2)
|
|
height = height1 + height2 - (endy - starty)
|
|
|
|
if width <= 0 or height <= 0:
|
|
ratio = 0
|
|
else:
|
|
Area = width * height
|
|
Area1 = width1 * height1
|
|
Area2 = width2 * height2
|
|
ratio = Area * 1. / (Area1 + Area2 - Area)
|
|
return ratio
|
|
|
|
|
|
parser = argparse.ArgumentParser(description='Package eval images')
|
|
# general
|
|
parser.add_argument('--data-dir', default='', help='')
|
|
parser.add_argument('--image-size', type=int, default=112, help='')
|
|
parser.add_argument('--gpu', type=int, default=0, help='')
|
|
parser.add_argument('--det-prefix', type=str, default='./model/R50', help='')
|
|
parser.add_argument('--output', default='./', help='path to save.')
|
|
parser.add_argument('--align-mode', default='arcface', help='align mode.')
|
|
args = parser.parse_args()
|
|
|
|
gpu_id = args.gpu
|
|
|
|
detector = RetinaFace(args.det_prefix, 0, gpu_id, network='net3')
|
|
target_size = 400
|
|
max_size = 800
|
|
|
|
|
|
def get_norm_crop(image_path):
|
|
im = cv2.imread(image_path)
|
|
im_shape = im.shape
|
|
im_size_min = np.min(im_shape[0:2])
|
|
im_size_max = np.max(im_shape[0:2])
|
|
im_scale = float(target_size) / float(im_size_min)
|
|
# prevent bigger axis from being more than max_size:
|
|
if np.round(im_scale * im_size_max) > max_size:
|
|
im_scale = float(max_size) / float(im_size_max)
|
|
bbox, landmark = detector.detect(im, threshold=0.5, scales=[im_scale])
|
|
#print(im.shape, bbox.shape, landmark.shape)
|
|
if bbox.shape[0] == 0:
|
|
bbox, landmark = detector.detect(
|
|
im,
|
|
threshold=0.05,
|
|
scales=[im_scale * 0.75, im_scale, im_scale * 2.0])
|
|
print('refine', im.shape, bbox.shape, landmark.shape)
|
|
nrof_faces = bbox.shape[0]
|
|
if nrof_faces > 0:
|
|
det = bbox[:, 0:4]
|
|
img_size = np.asarray(im.shape)[0:2]
|
|
bindex = 0
|
|
if nrof_faces > 1:
|
|
bounding_box_size = (det[:, 2] - det[:, 0]) * (det[:, 3] -
|
|
det[:, 1])
|
|
img_center = img_size / 2
|
|
offsets = np.vstack([(det[:, 0] + det[:, 2]) / 2 - img_center[1],
|
|
(det[:, 1] + det[:, 3]) / 2 - img_center[0]])
|
|
offset_dist_squared = np.sum(np.power(offsets, 2.0), 0)
|
|
bindex = np.argmax(bounding_box_size - offset_dist_squared *
|
|
2.0) # some extra weight on the centering
|
|
#_bbox = bounding_boxes[bindex, 0:4]
|
|
_landmark = landmark[bindex]
|
|
warped = face_align.norm_crop(im,
|
|
landmark=_landmark,
|
|
image_size=args.image_size,
|
|
mode=args.align_mode)
|
|
return warped
|
|
else:
|
|
return None
|
|
|
|
|
|
bins = []
|
|
issame_list = []
|
|
pp = 0
|
|
for line in open(os.path.join(args.data_dir, 'pairs_label.txt'), 'r'):
|
|
pp += 1
|
|
if pp % 100 == 0:
|
|
print('processing', pp)
|
|
line = line.strip().split()
|
|
assert len(line) == 3
|
|
path1 = os.path.join(args.data_dir, line[0])
|
|
path2 = os.path.join(args.data_dir, line[1])
|
|
im1 = get_norm_crop(path1)
|
|
im2 = get_norm_crop(path2)
|
|
issame = True
|
|
if line[2] == '0':
|
|
issame = False
|
|
issame_list.append(issame)
|
|
for im in [im1, im2]:
|
|
_, s = cv2.imencode('.jpg', im)
|
|
bins.append(s)
|
|
|
|
with open(args.output, 'wb') as f:
|
|
pickle.dump((bins, issame_list), f, protocol=pickle.HIGHEST_PROTOCOL)
|