From f0bf4e6cc962a0fb344d5ffbf1d07c9d890caaea Mon Sep 17 00:00:00 2001 From: Jia Guo Date: Tue, 8 Jan 2019 22:52:28 +0800 Subject: [PATCH] refine test for stacked dense unet --- alignment/SDUNet | 1 - alignment/alignment.py | 42 ---------------- alignment/benchmark.py | 65 ------------------------- alignment/draw.py | 70 -------------------------- alignment/test.py | 108 +++++++++++++++++------------------------ 5 files changed, 44 insertions(+), 242 deletions(-) delete mode 160000 alignment/SDUNet delete mode 100644 alignment/alignment.py delete mode 100644 alignment/benchmark.py delete mode 100644 alignment/draw.py mode change 100755 => 100644 alignment/test.py diff --git a/alignment/SDUNet b/alignment/SDUNet deleted file mode 160000 index 0e6060a..0000000 --- a/alignment/SDUNet +++ /dev/null @@ -1 +0,0 @@ -Subproject commit 0e6060a5a8db6a1761391a1ec7016fe31eb21732 diff --git a/alignment/alignment.py b/alignment/alignment.py deleted file mode 100644 index a2c264c..0000000 --- a/alignment/alignment.py +++ /dev/null @@ -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 - - - diff --git a/alignment/benchmark.py b/alignment/benchmark.py deleted file mode 100644 index 003675e..0000000 --- a/alignment/benchmark.py +++ /dev/null @@ -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) - - diff --git a/alignment/draw.py b/alignment/draw.py deleted file mode 100644 index fa689f7..0000000 --- a/alignment/draw.py +++ /dev/null @@ -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 diff --git a/alignment/test.py b/alignment/test.py old mode 100755 new mode 100644 index 1ef2d15..609bf2b --- a/alignment/test.py +++ b/alignment/test.py @@ -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)