Add image rotation for detecting more faces and dealing with awkward …

…angles (#253)

* Image rotator for extract and convert ready for testing

* Revert "Image rotator for extract and convert ready for testing"

This reverts commit bbeb19ef26d2392b0f5f954dd202137a70444d4e.

Error in extract code

* add image rotation support to detect more faces

* Update convert.py

Amended to do a single check for for rotation rather than checking twice. Performance gain is likely to be marginal to non-existent, but can't hurt.

* Update convert.py

remove type

* cli.py: Only output message on verbose. Convert.py: Only check for rotation amount once

* Changed command line flag to take arguments to ease future development

* Realigning for upstream/Master

* Minor fix

lib/cli.py

@@ -6,7 +6,7 @@
 from pathlib import Path
 from lib.FaceFilter import FaceFilter
 from lib.faces_detect import detect_faces, DetectedFace
-from lib.utils import get_image_paths, get_folder
+from lib.utils import get_image_paths, get_folder, rotate_image
 from lib import Serializer
 
 class FullPaths(argparse.Action):
@@ -142,6 +142,8 @@ def get_faces_alignments(self, filename, image):
         faces = self.faces_detected[os.path.basename(filename)]
         for rawface in faces:
             face = DetectedFace(**rawface)
+            # Rotate the image if necessary
+            if face.r != 0: image = rotate_image(image, face.r)
             face.image = image[face.y : face.y + face.h, face.x : face.x + face.w]
             if self.filter is not None and not self.filter.check(face):
                 print('Skipping not recognized face!')
@@ -154,9 +156,9 @@ def get_faces_alignments(self, filename, image):
             print('Note: Found more than one face in an image! File: %s' % filename)
             self.verify_output = True
 
-    def get_faces(self, image):
+    def get_faces(self, image, rotation=0):
         faces_count = 0
-        faces = detect_faces(image, self.arguments.detector)
+        faces = detect_faces(image, rotation, self.arguments.detector)
 
         for face in faces:
             if self.filter is not None and not self.filter.check(face):

lib/faces_detect.py

@@ -1,14 +1,15 @@
 from lib import FaceLandmarksExtractor
 
-def detect_faces(frame, model="hog"):
+def detect_faces(frame, rotation=0, model="hog"):
     fd = FaceLandmarksExtractor.extract (frame, True if model == "cnn" else False )
     for face in fd:
         x, y, right, bottom, landmarks = face[0][0], face[0][1], face[0][2], face[0][3], face[1]
-        yield DetectedFace(frame[y: bottom, x: right], x, right - x, y, bottom - y, landmarksXY=landmarks)
+        yield DetectedFace(frame[y: bottom, x: right], rotation, x, right - x, y, bottom - y, landmarksXY=landmarks)
 
 class DetectedFace(object):
-    def __init__(self, image=None, x=None, w=None, y=None, h=None, landmarksXY=None):
+    def __init__(self, image=None, r=None, x=None, w=None, y=None, h=None, landmarksXY=None):
         self.image = image
+        self.r = r
         self.x = x
         self.w = w
         self.y = y

lib/utils.py

@@ -1,3 +1,4 @@
+import cv2
 import sys
 from os.path import basename, exists
 
@@ -30,6 +31,20 @@ def get_image_paths(directory, exclude=[], debug=False):
 
     return dir_contents
 
+def rotate_image(image, angle):
+    ''' Rotates an image by 90, 180 or 270 degrees. Positive for clockwise, negative for 
+        counterclockwise '''
+    if angle < 0: angle = sum((360, angle))
+    if angle == 90:
+        image = cv2.flip(cv2.transpose(image),flipCode=1)
+    elif angle == 180:
+        image = cv2.flip(image,flipCode=-1)
+    elif angle == 270:
+        image = cv2.flip(cv2.transpose(image),flipCode=0)
+    else:
+        print('Unsupported image rotation angle: {}. Image unmodified'.format(angle))
+    return image
+
 # From: https://stackoverflow.com/questions/7323664/python-generator-pre-fetch
 import threading
 import queue as Queue

scripts/convert.py

@@ -6,7 +6,7 @@
 from tqdm import tqdm
 
 from lib.cli import DirectoryProcessor, FullPaths
-from lib.utils import BackgroundGenerator, get_folder, get_image_paths
+from lib.utils import BackgroundGenerator, get_folder, get_image_paths, rotate_image
 
 from plugins.PluginLoader import PluginLoader
 
@@ -226,8 +226,15 @@ def convert(self, converter, item):
                     if self.input_aligned_dir is not None and self.check_skipface(filename, idx):
                         print ('face {} for frame {} was deleted, skipping'.format(idx, os.path.basename(filename)))
                         continue
-                    image = converter.patch_image(image, face, 64 if "128" not in self.arguments.trainer else 128)
-                    # TODO: This switch between 64 and 128 is a hack for now. We should have a separate cli option for size
+                    # Check for image rotations and rotate before mapping face
+                    if face.r != 0:
+                        image = rotate_image(image, face.r)
+                        image = converter.patch_image(image, face, 64 if "128" not in self.arguments.trainer else 128)
+                        # TODO: This switch between 64 and 128 is a hack for now. We should have a separate cli option for size
+                        image = rotate_image(image, face.r * -1)
+                    else:
+                        image = converter.patch_image(image, face, 64 if "128" not in self.arguments.trainer else 128)
+                        # TODO: This switch between 64 and 128 is a hack for now. We should have a separate cli option for size
 
             output_file = get_folder(self.output_dir) / Path(filename).name
             cv2.imwrite(str(output_file), image)

scripts/extract.py

@@ -4,7 +4,7 @@
 from tqdm import tqdm
 import os
 
-from lib.cli import DirectoryProcessor
+from lib.cli import DirectoryProcessor, rotate_image
 from lib.utils import get_folder
 from lib.multithreading import pool_process
 from plugins.PluginLoader import PluginLoader
@@ -65,7 +65,15 @@ def add_optional_arguments(self, parser):
                             dest="debug_landmarks",
                             default=False,
                             help="Draw landmarks for debug.")
-
+
+        parser.add_argument('-r', '--rotate-images',
+                            type=str,
+                            dest="rotate_images",
+                            choices=("on", "off"),
+                            default="off",
+                            help="If a face isn't found, rotate the images through 90 degree "
+                                 "iterations to try to find a face. Can find more faces at the "
+                                 "cost of extraction speed.")
         return parser
 
     def process(self):
@@ -98,23 +106,40 @@ def processFiles(self, filename):
             pass
         return filename, []
 
+    def imageRotator(self, image):
+        ''' rotates the image through 90 degree iterations to find a face '''
+        angle = 90
+        while angle <= 270:
+            rotated_image = rotate_image(image, angle)
+            faces = self.get_faces(rotated_image, rotation=angle)
+            rotated_faces = [(idx, face) for idx, face in faces]
+            if len(rotated_faces) != 0:
+                if self.arguments.verbose:
+                    print('found face(s) by rotating image {} degrees'.format(angle))
+                break
+            angle += 90
+        return rotated_faces, rotated_image
+
     def handleImage(self, image, filename):
-        count = 0
-
         faces = self.get_faces(image)
+        process_faces = [(idx, face) for idx, face in faces]
+
+        # Run image rotator if requested and no faces found        
+        if self.arguments.rotate_images.lower() == 'on' and len(process_faces) == 0:
+            process_faces, image = self.imageRotator(image)
+
         rvals = []
-        for idx, face in faces:
-            count = idx
-
+        for idx, face in process_faces:
             # Draws landmarks for debug
             if self.arguments.debug_landmarks:
                 for (x, y) in face.landmarksAsXY():
                     cv2.circle(image, (x, y), 2, (0, 0, 255), -1)
-                    
+
             resized_image = self.extractor.extract(image, face, 256)
             output_file = get_folder(self.output_dir) / Path(filename).stem
             cv2.imwrite('{}_{}{}'.format(str(output_file), str(idx), Path(filename).suffix), resized_image)
             f = {
+                "r": face.r,
                 "x": face.x,
                 "w": face.w,
                 "y": face.y,