提交 e053c8c9 编写于 作者: B baiyfbupt

code clean

上级 de4504d6
......@@ -3,7 +3,7 @@ import time
import numpy as np
import argparse
import functools
import datetime
import cv2
from PIL import Image
from PIL import ImageDraw
......@@ -12,7 +12,6 @@ import paddle.fluid as fluid
import reader
from pyramidbox import PyramidBox
from utility import add_arguments, print_arguments
from paddle.fluid.framework import Program, Parameter, default_main_program, Variable
parser = argparse.ArgumentParser(description=__doc__)
add_arg = functools.partial(add_arguments, argparser=parser)
# yapf: disable
......@@ -21,8 +20,6 @@ add_arg('use_pyramidbox', bool, False, "Whether use PyramidBox model.")
add_arg('confs_threshold', float, 0.25, "Confidence threshold to draw bbox.")
add_arg('image_path', str, '', "The data root path.")
add_arg('model_dir', str, '', "The model path.")
add_arg('resize_h', int, 0, "The resized image height.")
add_arg('resize_w', int, 0, "The resized image height.")
# yapf: enable
......@@ -115,19 +112,7 @@ def bbox_vote(det):
return dets
def detect_face(image, image_shape, raw_image, shrink):
num_classes = 2
place = fluid.CUDAPlace(0) if args.use_gpu else fluid.CPUPlace()
exe = fluid.Executor(place)
if shrink != 1:
image = image.resize((int(image_shape[2] * shrink),
int(image_shape[1] * shrink)), Image.ANTIALIAS)
image_shape = [
image_shape[0], int(image_shape[1] * shrink),
int(image_shape[2] * shrink)
]
print "image_shape:", image_shape
def image_preprocess(image):
img = np.array(image)
# HWC to CHW
if len(img.shape) == 3:
......@@ -141,47 +126,54 @@ def detect_face(image, image_shape, raw_image, shrink):
img = img * 0.007843
img = [img]
img = np.array(img)
return img
def detect_face(image, shrink):
image_shape = [3, image.size[1], image.size[0]]
num_classes = 2
place = fluid.CUDAPlace(0) if args.use_gpu else fluid.CPUPlace()
exe = fluid.Executor(place)
if shrink != 1:
image = image.resize((int(image_shape[2] * shrink),
int(image_shape[1] * shrink)), Image.ANTIALIAS)
image_shape = [
image_shape[0], int(image_shape[1] * shrink),
int(image_shape[2] * shrink)
]
print "image_shape:", image_shape
img = image_preprocess(image)
scope = fluid.core.Scope()
model_program = fluid.Program()
main_program = fluid.Program()
startup_program = fluid.Program()
with fluid.scope_guard(scope):
with fluid.unique_name.guard():
with fluid.program_guard(model_program, startup_program):
with fluid.program_guard(main_program, startup_program):
fetches = []
network = PyramidBox(
image_shape,
num_classes,
sub_network=args.use_pyramidbox,
is_infer=True)
infer_program, nmsed_out = network.infer()
infer_program, nmsed_out = network.infer(main_program)
fetches = [nmsed_out]
feeder = fluid.DataFeeder(
place=place, feed_list=network.feeds())
fluid.io.load_persistables(
exe, args.model_dir, main_program=model_program)
#fluid.io.load_vars(exe, args.model_dir, predicate=if_exist)
exe, args.model_dir, main_program=main_program)
detection, = exe.run(infer_program,
feed=feeder.feed([img]),
feed={'image': img},
fetch_list=fetches,
return_numpy=False)
detection = np.array(detection)
# layout: xmin, ymin, xmax. ymax, score
det_conf = detection[:, 1]
if args.resize_h != 0 and args.resize_w != 0:
det_xmin = raw_image.size[0] * detection[:, 2]
det_ymin = raw_image.size[1] * detection[:, 3]
det_xmax = raw_image.size[0] * detection[:, 4]
det_ymax = raw_image.size[1] * detection[:, 5]
else:
det_xmin = image_shape[2] * detection[:, 2] / shrink
det_ymin = image_shape[1] * detection[:, 3] / shrink
det_xmax = image_shape[2] * detection[:, 4] / shrink
det_ymax = image_shape[1] * detection[:, 5] / shrink
det_xmin = image_shape[2] * detection[:, 2] / shrink
det_ymin = image_shape[1] * detection[:, 3] / shrink
det_xmax = image_shape[2] * detection[:, 4] / shrink
det_ymax = image_shape[1] * detection[:, 5] / shrink
det = np.column_stack((det_xmin, det_ymin, det_xmax, det_ymax, det_conf))
keep_index = np.where(det[:, 4] >= 0)[0]
......@@ -189,40 +181,37 @@ def detect_face(image, image_shape, raw_image, shrink):
return det
def flip_test(image, image_shape, raw_image, shrink):
def flip_test(image, shrink):
image = image.transpose(Image.FLIP_LEFT_RIGHT)
det_f = detect_face(image, image_shape, raw_image, shrink)
det_f = detect_face(image, shrink)
det_t = np.zeros(det_f.shape)
det_t[:, 0] = raw_image.size[0] - det_f[:, 2]
det_t[:, 0] = image.size[0] - det_f[:, 2]
det_t[:, 1] = det_f[:, 1]
det_t[:, 2] = raw_image.size[0] - det_f[:, 0]
det_t[:, 2] = image.size[0] - det_f[:, 0]
det_t[:, 3] = det_f[:, 3]
det_t[:, 4] = det_f[:, 4]
return det_t
def multi_scale_test(image, image_shape, raw_image, max_im_shrink):
def multi_scale_test(image, max_shrink):
# shrink detecting and shrink only detect big face
st = 0.5 if max_im_shrink >= 0.75 else 0.5 * max_im_shrink
det_s = detect_face(image, image_shape, raw_image, st)
st = 0.5 if max_shrink >= 0.75 else 0.5 * max_shrink
det_s = detect_face(image, st)
index = np.where(
np.maximum(det_s[:, 2] - det_s[:, 0] + 1, det_s[:, 3] - det_s[:, 1] + 1)
> 30)[0]
det_s = det_s[index, :]
# enlarge one times
bt = min(2, max_im_shrink) if max_im_shrink > 1 else (
st + max_im_shrink) / 2
det_b = detect_face(image, image_shape, raw_image, bt)
bt = min(2, max_shrink) if max_shrink > 1 else (st + max_shrink) / 2
det_b = detect_face(image, bt)
# enlarge small iamge x times for small face
if max_im_shrink > 2:
# enlarge small image x times for small face
if max_shrink > 2:
bt *= 2
while bt < max_im_shrink:
det_b = np.row_stack(
(det_b, detect_face(image, image_shape, raw_image, bt)))
while bt < max_shrink:
det_b = np.row_stack((det_b, detect_face(image, bt)))
bt *= 2
det_b = np.row_stack(
(det_b, detect_face(image, image_shape, raw_image, max_im_shrink)))
det_b = np.row_stack((det_b, detect_face(image, max_shrink)))
# enlarge only detect small face
if bt > 1:
......@@ -239,30 +228,28 @@ def multi_scale_test(image, image_shape, raw_image, max_im_shrink):
def get_im_shrink(image_shape):
max_im_shrink_v1 = (0x7fffffff / 577.0 /
(image_shape[1] * image_shape[2]))**0.5
max_im_shrink_v2 = (
max_shrink_v1 = (0x7fffffff / 577.0 /
(image_shape[1] * image_shape[2]))**0.5
max_shrink_v2 = (
(678 * 1024 * 2.0 * 2.0) / (image_shape[1] * image_shape[2]))**0.5
max_im_shrink = get_round(min(max_im_shrink_v1, max_im_shrink_v2), 2) - 0.3
if max_im_shrink >= 1.5 and max_im_shrink < 2:
max_im_shrink = max_im_shrink - 0.1
elif max_im_shrink >= 2 and max_im_shrink < 3:
max_im_shrink = max_im_shrink - 0.2
elif max_im_shrink >= 3 and max_im_shrink < 4:
max_im_shrink = max_im_shrink - 0.3
elif max_im_shrink >= 4 and max_im_shrink < 5:
max_im_shrink = max_im_shrink - 0.4
elif max_im_shrink >= 5 and max_im_shrink < 6:
max_im_shrink = max_im_shrink - 0.5
elif max_im_shrink >= 6:
max_im_shrink = max_im_shrink - 0.5
print 'max_im_shrink = ', max_im_shrink
shrink = max_im_shrink if max_im_shrink < 1 else 1
max_shrink = get_round(min(max_shrink_v1, max_shrink_v2), 2) - 0.3
if max_shrink >= 1.5 and max_shrink < 2:
max_shrink = max_shrink - 0.1
elif max_shrink >= 2 and max_shrink < 3:
max_shrink = max_shrink - 0.2
elif max_shrink >= 3 and max_shrink < 4:
max_shrink = max_shrink - 0.3
elif max_shrink >= 4 and max_shrink < 5:
max_shrink = max_shrink - 0.4
elif max_shrink >= 5:
max_shrink = max_shrink - 0.5
print 'max_shrink = ', max_shrink
shrink = max_shrink if max_shrink < 1 else 1
print "shrink = ", shrink
return shrink, max_im_shrink
return shrink, max_shrink
def infer(args, batch_size, data_args):
......@@ -277,33 +264,25 @@ def infer(args, batch_size, data_args):
image = img[0][0]
image_path = img[0][1]
raw_image = Image.open(image_path)
image_shape = [3, image.size[1], image.size[0]]
if args.resize_h != 0 and args.resize_w != 0:
image_shape = [3, args.resize_h, args.resize_w]
else:
image_shape = [3, image.size[1], image.size[0]]
shrink, max_im_shrink = get_im_shrink(image_shape)
det0 = detect_face(image, image_shape, raw_image, shrink)
det1 = flip_test(image, image_shape, raw_image, shrink)
[det2, det3] = multi_scale_test(image, image_shape, raw_image,
max_im_shrink)
shrink, max_shrink = get_im_shrink(image_shape)
det0 = detect_face(image, shrink)
det1 = flip_test(image, shrink)
[det2, det3] = multi_scale_test(image, max_shrink)
det = np.row_stack((det0, det1, det2, det3))
dets = bbox_vote(det)
image_name = image_path.split('/')[-1]
image_class = image_path.split('/')[-2]
if not os.path.exists('./infer_results/' + image_class.encode('utf-8')):
os.makedirs('./infer_results/' + image_class.encode('utf-8'))
f = open('./infer_results/' + image_class.encode('utf-8') + '/' +
image_name.encode('utf-8')[:-4] + '.txt', 'w')
write_to_txt(image_path, f, dets)
#draw_bounding_box_on_image(image_path, dets, args.confs_threshold)
# write_to_txt(image_path, f, dets)
# draw_bounding_box_on_image(image_path, dets, args.confs_threshold)
print "Done"
......@@ -316,8 +295,6 @@ if __name__ == '__main__':
data_args = reader.Settings(
data_dir=data_dir,
resize_h=args.resize_h,
resize_w=args.resize_w,
mean_value=[104., 117., 123],
apply_distort=False,
apply_expand=False,
......
......@@ -269,9 +269,6 @@ def pyramidbox(settings, file_list, mode, shuffle):
yield im, boxes, expand_bboxes(boxes), lbls, difficults
if mode == 'test':
if settings.resize_w and settings.resize_h:
im = im.resize((settings.resize_w, settings.resize_h),
Image.ANTIALIAS)
yield im, image_path
return reader
......
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册