mirror of
https://github.com/deepinsight/insightface.git
synced 2026-05-17 14:26:08 +00:00
refine test for stacked dense unet
This commit is contained in:
Jia Guo
Submodule alignment/SDUNet deleted from 0e6060a5a8
@@ -1,42 +0,0 @@
|
||||
import argparse
|
||||
import cv2
|
||||
import numpy as np
|
||||
import sys
|
||||
import mxnet as mx
|
||||
import datetime
|
||||
|
||||
class Alignment:
|
||||
def __init__(self, prefix, epoch, ctx_id=0):
|
||||
print('loading',prefix, epoch)
|
||||
ctx = mx.gpu(ctx_id)
|
||||
sym, arg_params, aux_params = mx.model.load_checkpoint(prefix, epoch)
|
||||
all_layers = sym.get_internals()
|
||||
sym = all_layers['heatmap_output']
|
||||
image_size = (128, 128)
|
||||
self.image_size = image_size
|
||||
model = mx.mod.Module(symbol=sym, context=ctx, label_names = None)
|
||||
#model = mx.mod.Module(symbol=sym, context=ctx)
|
||||
model.bind(for_training=False, data_shapes=[('data', (1, 3, image_size[0], image_size[1]))])
|
||||
model.set_params(arg_params, aux_params)
|
||||
self.model = model
|
||||
|
||||
def get(self, img):
|
||||
rimg = cv2.resize(img, (self.image_size[1], self.image_size[0]))
|
||||
img = cv2.cvtColor(rimg, cv2.COLOR_BGR2RGB)
|
||||
img = np.transpose(img, (2,0,1)) #3*112*112, RGB
|
||||
input_blob = np.zeros( (1, 3, self.image_size[1], self.image_size[0]),dtype=np.uint8 )
|
||||
input_blob[0] = img
|
||||
data = mx.nd.array(input_blob)
|
||||
db = mx.io.DataBatch(data=(data,))
|
||||
self.model.forward(db, is_train=False)
|
||||
alabel = self.model.get_outputs()[-1].asnumpy()[0]
|
||||
ret = np.zeros( (alabel.shape[0], 2), dtype=np.float32)
|
||||
for i in xrange(alabel.shape[0]):
|
||||
a = cv2.resize(alabel[i], (self.image_size[1], self.image_size[0]))
|
||||
ind = np.unravel_index(np.argmax(a, axis=None), a.shape)
|
||||
#ret[i] = (ind[0], ind[1]) #h, w
|
||||
ret[i] = (ind[1], ind[0]) #w, h
|
||||
return ret
|
||||
|
||||
|
||||
|
||||
@@ -1,65 +0,0 @@
|
||||
import argparse
|
||||
import cv2
|
||||
import numpy as np
|
||||
import sys
|
||||
import mxnet as mx
|
||||
import datetime
|
||||
|
||||
parser = argparse.ArgumentParser(description='face model test')
|
||||
# general
|
||||
parser.add_argument('--image-size', default='128,128', help='')
|
||||
parser.add_argument('--model', default='./models/test,15', help='path to load model.')
|
||||
parser.add_argument('--gpu', default=0, type=int, help='gpu id')
|
||||
parser.add_argument('--batch-size', default=10, type=int, help='batch size')
|
||||
parser.add_argument('--iterations', default=10, type=int, help='iterations')
|
||||
args = parser.parse_args()
|
||||
|
||||
_vec = args.image_size.split(',')
|
||||
assert len(_vec)==2
|
||||
image_size = (int(_vec[0]), int(_vec[1]))
|
||||
_vec = args.model.split(',')
|
||||
assert len(_vec)==2
|
||||
prefix = _vec[0]
|
||||
epoch = int(_vec[1])
|
||||
print('loading',prefix, epoch)
|
||||
if args.gpu>=0:
|
||||
ctx = mx.gpu(args.gpu)
|
||||
else:
|
||||
ctx = mx.cpu()
|
||||
sym, arg_params, aux_params = mx.model.load_checkpoint(prefix, epoch)
|
||||
all_layers = sym.get_internals()
|
||||
sym = all_layers['heatmap_output']
|
||||
model = mx.mod.Module(symbol=sym, context=ctx, label_names = None)
|
||||
#model = mx.mod.Module(symbol=sym, context=ctx)
|
||||
model.bind(for_training=False, data_shapes=[('data', (args.batch_size, 3, image_size[0], image_size[1]))])
|
||||
#model.bind(for_training=False, data_shapes=[('data', (args.batch_size, 3, image_size[0], image_size[1]))], label_shapes=[('softmax_label', (args.batch_size,84,64,64))])
|
||||
model.set_params(arg_params, aux_params)
|
||||
img_path = './test.png'
|
||||
|
||||
img = cv2.imread(img_path)
|
||||
|
||||
rimg = cv2.resize(img, (image_size[1], image_size[0]))
|
||||
img = cv2.cvtColor(rimg, cv2.COLOR_BGR2RGB)
|
||||
img = np.transpose(img, (2,0,1)) #3*112*112, RGB
|
||||
input_blob = np.zeros( (args.batch_size, 3, image_size[1], image_size[0]),dtype=np.uint8 )
|
||||
for i in xrange(args.batch_size):
|
||||
input_blob[i] = img
|
||||
data = mx.nd.array(input_blob)
|
||||
print(data.shape)
|
||||
label = mx.nd.zeros( (args.batch_size, 84, 64, 64) )
|
||||
#db = mx.io.DataBatch(data=(data,))
|
||||
db = mx.io.DataBatch(data=(data,), label=(label,))
|
||||
stat = []
|
||||
warmup = 2
|
||||
for i in xrange(args.iterations+warmup):
|
||||
#print(i)
|
||||
time_now = datetime.datetime.now()
|
||||
model.forward(db, is_train=False)
|
||||
output = model.get_outputs()[-1].asnumpy()
|
||||
time_now2 = datetime.datetime.now()
|
||||
diff = time_now2 - time_now
|
||||
stat.append(diff.total_seconds())
|
||||
stat = stat[warmup:]
|
||||
print(np.mean(stat)/args.batch_size)
|
||||
|
||||
|
||||
@@ -1,70 +0,0 @@
|
||||
import numpy as np
|
||||
import skimage.draw
|
||||
|
||||
def line(img, pt1, pt2, color, width):
|
||||
# Draw a line on an image
|
||||
# Make sure dimension of color matches number of channels in img
|
||||
|
||||
# First get coordinates for corners of the line
|
||||
diff = np.array([pt1[1] - pt2[1], pt1[0] - pt2[0]], np.float)
|
||||
mag = np.linalg.norm(diff)
|
||||
if mag >= 1:
|
||||
diff *= width / (2 * mag)
|
||||
x = np.array([pt1[0] - diff[0], pt2[0] - diff[0], pt2[0] + diff[0], pt1[0] + diff[0]], int)
|
||||
y = np.array([pt1[1] + diff[1], pt2[1] + diff[1], pt2[1] - diff[1], pt1[1] - diff[1]], int)
|
||||
else:
|
||||
d = float(width) / 2
|
||||
x = np.array([pt1[0] - d, pt1[0] + d, pt1[0] + d, pt1[0] - d], int)
|
||||
y = np.array([pt1[1] - d, pt1[1] - d, pt1[1] + d, pt1[1] + d], int)
|
||||
|
||||
# noinspection PyArgumentList
|
||||
rr, cc = skimage.draw.polygon(y, x, img.shape)
|
||||
img[rr, cc] = color
|
||||
|
||||
return img
|
||||
|
||||
def limb(img, pt1, pt2, color, width):
|
||||
# Specific handling of a limb, in case the annotation isn't there for one of the joints
|
||||
if pt1[0] > 0 and pt2[0] > 0:
|
||||
line(img, pt1, pt2, color, width)
|
||||
elif pt1[0] > 0:
|
||||
circle(img, pt1, color, width)
|
||||
elif pt2[0] > 0:
|
||||
circle(img, pt2, color, width)
|
||||
|
||||
def gaussian(img, pt, sigma):
|
||||
# Draw a 2D gaussian
|
||||
|
||||
# Check that any part of the gaussian is in-bounds
|
||||
ul = [int(pt[0] - 3 * sigma), int(pt[1] - 3 * sigma)]
|
||||
br = [int(pt[0] + 3 * sigma + 1), int(pt[1] + 3 * sigma + 1)]
|
||||
if (ul[0] > img.shape[1] or ul[1] >= img.shape[0] or
|
||||
br[0] < 0 or br[1] < 0):
|
||||
# If not, just return the image as is
|
||||
return img
|
||||
|
||||
# Generate gaussian
|
||||
size = 6 * sigma + 1
|
||||
x = np.arange(0, size, 1, float)
|
||||
y = x[:, np.newaxis]
|
||||
x0 = y0 = size // 2
|
||||
# The gaussian is not normalized, we want the center value to equal 1
|
||||
g = np.exp(- ((x - x0) ** 2 + (y - y0) ** 2) / (2 * sigma ** 2))
|
||||
|
||||
# Usable gaussian range
|
||||
g_x = max(0, -ul[0]), min(br[0], img.shape[1]) - ul[0]
|
||||
g_y = max(0, -ul[1]), min(br[1], img.shape[0]) - ul[1]
|
||||
# Image range
|
||||
img_x = max(0, ul[0]), min(br[0], img.shape[1])
|
||||
img_y = max(0, ul[1]), min(br[1], img.shape[0])
|
||||
|
||||
img[img_y[0]:img_y[1], img_x[0]:img_x[1]] = g[g_y[0]:g_y[1], g_x[0]:g_x[1]]
|
||||
return img
|
||||
|
||||
def circle(img, pt, color, radius):
|
||||
# Draw a circle
|
||||
# Mostly a convenient wrapper for skimage.draw.circle
|
||||
|
||||
rr, cc = skimage.draw.circle(pt[1], pt[0], radius, img.shape)
|
||||
img[rr, cc] = color
|
||||
return img
|
||||
108
alignment/test.py
Executable file → Normal file
108
alignment/test.py
Executable file → Normal file
@@ -1,71 +1,51 @@
|
||||
import argparse
|
||||
import cv2
|
||||
import sys
|
||||
import numpy as np
|
||||
import datetime
|
||||
from alignment import Alignment
|
||||
sys.path.append('../SSH')
|
||||
from ssh_detector import SSHDetector
|
||||
import os
|
||||
import mxnet as mx
|
||||
|
||||
#short_max = 800
|
||||
scales = [1200, 1600]
|
||||
t = 2
|
||||
|
||||
detector = SSHDetector('../SSH/model/e2ef', 0)
|
||||
alignment = Alignment('./model/3d_I5', 12)
|
||||
out_filename = './out.png'
|
||||
class Handler:
|
||||
def __init__(self, prefix, epoch, ctx_id=0):
|
||||
print('loading',prefix, epoch)
|
||||
ctx = mx.gpu(ctx_id)
|
||||
sym, arg_params, aux_params = mx.model.load_checkpoint(prefix, epoch)
|
||||
all_layers = sym.get_internals()
|
||||
sym = all_layers['heatmap_output']
|
||||
image_size = (128, 128)
|
||||
self.image_size = image_size
|
||||
model = mx.mod.Module(symbol=sym, context=ctx, label_names = None)
|
||||
#model = mx.mod.Module(symbol=sym, context=ctx)
|
||||
model.bind(for_training=False, data_shapes=[('data', (1, 3, image_size[0], image_size[1]))])
|
||||
model.set_params(arg_params, aux_params)
|
||||
self.model = model
|
||||
|
||||
def get(self, img):
|
||||
rimg = cv2.resize(img, (self.image_size[1], self.image_size[0]))
|
||||
img = cv2.cvtColor(rimg, cv2.COLOR_BGR2RGB)
|
||||
img = np.transpose(img, (2,0,1)) #3*112*112, RGB
|
||||
input_blob = np.zeros( (1, 3, self.image_size[1], self.image_size[0]),dtype=np.uint8 )
|
||||
input_blob[0] = img
|
||||
data = mx.nd.array(input_blob)
|
||||
db = mx.io.DataBatch(data=(data,))
|
||||
self.model.forward(db, is_train=False)
|
||||
alabel = self.model.get_outputs()[-1].asnumpy()[0]
|
||||
ret = np.zeros( (alabel.shape[0], 2), dtype=np.float32)
|
||||
for i in xrange(alabel.shape[0]):
|
||||
a = cv2.resize(alabel[i], (self.image_size[1], self.image_size[0]))
|
||||
ind = np.unravel_index(np.argmax(a, axis=None), a.shape)
|
||||
#ret[i] = (ind[0], ind[1]) #h, w
|
||||
ret[i] = (ind[1], ind[0]) #w, h
|
||||
return ret
|
||||
|
||||
ctx_id = 0
|
||||
img_path = './test.png'
|
||||
img = cv2.imread(img_path)
|
||||
|
||||
handler = Handler('./model/SDU', 1, ctx_id)
|
||||
landmark = handler.get(img)
|
||||
#visualize landmark
|
||||
|
||||
|
||||
f = '../sample-images/t1.jpg'
|
||||
if len(sys.argv)>1:
|
||||
f = sys.argv[1]
|
||||
img = cv2.imread(f)
|
||||
im_shape = img.shape
|
||||
print(im_shape)
|
||||
target_size = scales[0]
|
||||
max_size = scales[1]
|
||||
im_size_min = np.min(im_shape[0:2])
|
||||
im_size_max = np.max(im_shape[0:2])
|
||||
if im_size_min>target_size or im_size_max>max_size:
|
||||
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)
|
||||
img = cv2.resize(img, None, None, fx=im_scale, fy=im_scale)
|
||||
print('resize to', img.shape)
|
||||
for i in xrange(t-1): #warmup
|
||||
faces = detector.detect(img, 0.5)
|
||||
timea = datetime.datetime.now()
|
||||
faces = detector.detect(img, 0.5)
|
||||
timeb = datetime.datetime.now()
|
||||
diff = timeb - timea
|
||||
print('detection uses', diff.total_seconds(), 'seconds')
|
||||
print('find', faces.shape[0], 'faces')
|
||||
|
||||
for face in faces:
|
||||
#print(face)
|
||||
cv2.rectangle(img, (face[0], face[1]), (face[2], face[3]), (255, 0, 0), 1)
|
||||
w = face[2] - face[0]
|
||||
h = face[3] - face[1]
|
||||
wc = int( (face[2]+face[0])/2 )
|
||||
hc = int( (face[3]+face[1])/2 )
|
||||
size = int(max(w, h)*1.3)
|
||||
scale = 100.0/max(w,h)
|
||||
M = [
|
||||
[scale, 0, 64-wc*scale],
|
||||
[0, scale, 64-hc*scale],
|
||||
]
|
||||
M = np.array(M)
|
||||
IM = cv2.invertAffineTransform(M)
|
||||
#print(M, IM)
|
||||
ebox = cv2.warpAffine(img, M, (128, 128))
|
||||
#ebox = cv2.getRectSubPix(img, (size, size), (wc, hc))
|
||||
landmark = alignment.get(ebox)
|
||||
#print(landmark.shape)
|
||||
for l in range(landmark.shape[0]):
|
||||
point = np.ones( (3,), dtype=np.float32)
|
||||
point[0:2] = landmark[l]
|
||||
point = np.dot(IM, point)
|
||||
pp = (int(point[0]), int(point[1]))
|
||||
#print(pp)
|
||||
cv2.circle(img, (pp[0], pp[1]), 1, (0, 0, 255), 1)
|
||||
print('write to', out_filename)
|
||||
cv2.imwrite(out_filename, img)
|
||||
|
||||
Reference in New Issue
Block a user