davidsandberg / facenet

Face recognition using Tensorflow
MIT License
13.77k stars 4.81k forks source link

How to extract multiple faces in align_dataset_mtcnn.py #411

Closed Victoria2333 closed 7 years ago

Victoria2333 commented 7 years ago

I execute: python src/align/align_dataset_mtcnn.py input output --image_size 160 --margin 32 --random_order And there are multiple faces in one of those images,but the result only shows one face in each image,how can i modify this code? Please give me some tips~~~

moonhoenlee commented 7 years ago

I wrote down the following code to detect multiple faces. You can use the "--show_multiple_faces" option when using it.

@@ -36,6 +36,8 @@ import align.detect_face
 import random
 from time import sleep

+import cv2
+
 def main(args):
     sleep(random.random())
     output_dir = os.path.expanduser(args.output_dir)
@@ -96,26 +98,54 @@ def main(args):
                         bounding_boxes, _ = align.detect_face.detect_face(img, minsize, pnet, rnet, onet, threshold, factor)
                         nrof_faces = bounding_boxes.shape[0]
                         if nrof_faces>0:
-                            det = bounding_boxes[:,0:4]
-                            img_size = np.asarray(img.shape)[0:2]
-                            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)
-                                index = np.argmax(bounding_box_size-offset_dist_squared*2.0) # some extra weight on the centering
-                                det = det[index,:]
-                            det = np.squeeze(det)
-                            bb = np.zeros(4, dtype=np.int32)
-                            bb[0] = np.maximum(det[0]-args.margin/2, 0)
-                            bb[1] = np.maximum(det[1]-args.margin/2, 0)
-                            bb[2] = np.minimum(det[2]+args.margin/2, img_size[1])
-                            bb[3] = np.minimum(det[3]+args.margin/2, img_size[0])
-                            cropped = img[bb[1]:bb[3],bb[0]:bb[2],:]
-                            scaled = misc.imresize(cropped, (args.image_size, args.image_size), interp='bilinear')
-                            nrof_successfully_aligned += 1
-                            misc.imsave(output_filename, scaled)
-                            text_file.write('%s %d %d %d %d\n' % (output_filename, bb[0], bb[1], bb[2], bb[3]))
+                            if args.generate_multiple_faces:
+                                img_size = np.asarray(img.shape)[0:2]
+                                if args.show_multiple_faces:
+                                    cv_img = cv2.imread(image_path, cv2.IMREAD_COLOR)
+                                for i in range(nrof_faces):
+                                    bb = np.zeros(4, dtype=np.int32)
+                                    bb[0] = np.maximum(bounding_boxes[i][0]-args.margin/2, 0)
+                                    bb[1] = np.maximum(bounding_boxes[i][1]-args.margin/2, 0)
+                                    bb[2] = np.minimum(bounding_boxes[i][2]+args.margin/2, img_size[1])
+                                    bb[3] = np.minimum(bounding_boxes[i][3]+args.margin/2, img_size[0])
+                                    if args.show_multiple_faces:
+                                        cv2.rectangle(cv_img, (int(bb[0]),
+                                            int(bb[1])), (int(bb[2]),
+                                            int(bb[3])), (0, 255, 0), 2)
+
+                                    cropped = img[bb[1]:bb[3],bb[0]:bb[2],:]
+                                    scaled = misc.imresize(cropped, (args.image_size, args.image_size), interp='bilinear')
+                                    saveimage_filename = os.path.join(output_class_dir,
+                                            filename + "_%d.png" % i)
+                                    misc.imsave(saveimage_filename, scaled)
+                                    text_file.write('%s %d %d %d %d\n' % (saveimage_filename, bb[0], bb[1], bb[2], bb[3]))
+                                if args.show_multiple_faces:
+                                    cv2.imshow('image', cv_img)
+                                    cv2.waitKey(0)
+                                    cv2.destroyAllWindows()
+                                nrof_successfully_aligned += 1
+
+                            else:
+                                det = bounding_boxes[:,0:4]
+                                img_size = np.asarray(img.shape)[0:2]
+                                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)
+                                    index = np.argmax(bounding_box_size-offset_dist_squared*2.0) # some extra weight on the centering
+                                    det = det[index,:]
+                                det = np.squeeze(det)
+                                bb = np.zeros(4, dtype=np.int32)
+                                bb[0] = np.maximum(det[0]-args.margin/2, 0)
+                                bb[1] = np.maximum(det[1]-args.margin/2, 0)
+                                bb[2] = np.minimum(det[2]+args.margin/2, img_size[1])
+                                bb[3] = np.minimum(det[3]+args.margin/2, img_size[0])
+                                cropped = img[bb[1]:bb[3],bb[0]:bb[2],:]
+                                scaled = misc.imresize(cropped, (args.image_size, args.image_size), interp='bilinear')
+                                nrof_successfully_aligned += 1
+                                misc.imsave(output_filename, scaled)
+                                text_file.write('%s %d %d %d %d\n' % (output_filename, bb[0], bb[1], bb[2], bb[3]))
                         else:
                             print('Unable to align "%s"' % image_path)
                             text_file.write('%s\n' % (output_filename))
@@ -137,6 +167,10 @@ def parse_arguments(argv):
         help='Shuffles the order of images to enable alignment using multiple processes.', action='store_true')
     parser.add_argument('--gpu_memory_fraction', type=float,
         help='Upper bound on the amount of GPU memory that will be used by the process.', default=1.0)
+    parser.add_argument('--generate_multiple_faces', help='Generate multiple faces',
+            action='store_true')
+    parser.add_argument('--show_multiple_faces', help='Show multiple faces',
+            action='store_true')
     return parser.parse_args(argv)

 if __name__ == '__main__':
Victoria2333 commented 6 years ago

thank you,I already modify it by myself~