Adds support for arbitrary image rotations (#309)

* Add support for user-specified rotation angle in extract

* Added rotation-angle-list option to enumerate a list of angles to rotate through

* Adjust rotation matrix translation coords to avoid cropping

* Merged rotation-angle and rotation-angle-list options into rotate_images option

* Backwards compatibility

* Updated check whether to run image rotator

* Switched rotation convention to use positive angle = clockwise rotation, for backwards compatibility

lib/cli.py

@@ -41,6 +41,7 @@ class DirectoryProcessor(object):
     num_faces_detected = 0
     faces_detected = dict()
     verify_output = False
+    rotation_angles = None
 
     def __init__(self, subparser, command, description='default'):
         self.create_parser(subparser, command, description)
@@ -60,8 +61,22 @@ def process_arguments(self, arguments):
             self.serializer = Serializer.get_serializer(self.arguments.serializer or "json")
         print("Using {} serializer".format(self.serializer.ext))
 
+        try:
+            if self.arguments.rotate_images is not None and self.arguments.rotate_images != "off":
+                if self.arguments.rotate_images == "on":
+                    self.rotation_angles = range(90, 360, 90)
+                else:
+                    rotation_angles = [int(angle) for angle in self.arguments.rotate_images.split(",")]
+                    if len(rotation_angles) == 1:
+                        rotation_step_size = rotation_angles[0]
+                        self.rotation_angles = range(rotation_step_size, 360, rotation_step_size)
+                    elif len(rotation_angles) > 1:
+                        self.rotation_angles = rotation_angles
+        except AttributeError:
+            pass
+
         print('Starting, this may take a while...')
-        
+
         try:
             if self.arguments.skip_existing:
                 self.already_processed = get_image_paths(self.arguments.output_dir)

lib/utils.py

@@ -31,19 +31,21 @@ 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/22041699/rotate-an-image-without-cropping-in-opencv-in-c
+def rotate_image(image, angle, rotated_width=None, rotated_height=None):
+    height, width = image.shape[:2]
+    image_center = (width/2, height/2)
+    rotation_matrix = cv2.getRotationMatrix2D(image_center, -1.*angle, 1.)
+    if rotated_width is None or rotated_height is None:
+        abs_cos = abs(rotation_matrix[0, 0])
+        abs_sin = abs(rotation_matrix[0, 1])
+        if rotated_width is None:
+            rotated_width = int(height*abs_sin + width*abs_cos)
+        if rotated_height is None:
+            rotated_height = int(height*abs_cos + width*abs_sin)
+    rotation_matrix[0, 2] += rotated_width/2 - image_center[0]
+    rotation_matrix[1, 2] += rotated_height/2 - image_center[1]
+    return cv2.warpAffine(image, rotation_matrix, (rotated_width, rotated_height))
 
 # From: https://stackoverflow.com/questions/7323664/python-generator-pre-fetch
 import threading

scripts/convert.py

@@ -242,10 +242,11 @@ def convert(self, converter, item):
                         continue
                     # Check for image rotations and rotate before mapping face
                     if face.r != 0:
+                        height, width = image.shape[:2]
                         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)
+                        image = rotate_image(image, face.r * -1, rotated_width=width, rotated_height=height)
                     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

scripts/extract.py

@@ -71,11 +71,10 @@ def add_optional_arguments(self, parser):
         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.")
+                            default=None,
+                            help="If a face isn't found, rotate the images to try to find a face. Can find more faces at the "
+                                 "cost of extraction speed.  Pass in a single number to use increments of that size up to 360, "
+                                 "or pass in a list of numbers to enumerate exactly what angles to check.")
 
         parser.add_argument('-ae', '--align-eyes',
                             action="store_true",
@@ -123,26 +122,28 @@ def processFiles(self, filename):
             pass
         return filename, []
 
+    def getRotatedImageFaces(self, image, angle):
+        rotated_image = rotate_image(image, angle)
+        faces = self.get_faces(rotated_image, rotation=angle)
+        rotated_faces = [(idx, face) for idx, face in faces]
+        return rotated_faces, rotated_image
+
     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:
+        ''' rotates the image through rotation_angles to try to find a face '''
+        for angle in self.rotation_angles:
+            rotated_faces, rotated_image = self.getRotatedImageFaces(image, angle)
+            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):
         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:
+        if self.rotation_angles is not None and len(process_faces) == 0:
             process_faces, image = self.imageRotator(image)
 
         rvals = []