code clean

fluid/face_detection/infer.py

@@ -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,

fluid/face_detection/reader.py

@@ -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