Added ability to skip images with more than one face; fixed threshold type

This commit is contained in:
Filipp Bakanov 2016-05-03 22:15:57 +03:00
parent 5b46a5fcb4
commit d18017a195
3 changed files with 13 additions and 6 deletions

View File

@ -108,19 +108,21 @@ class AlignDlib:
# In rare cases, exceptions are thrown.
return []
def getLargestFaceBoundingBox(self, rgbImg):
def getLargestFaceBoundingBox(self, rgbImg, skipMulti = False):
"""
Find the largest face bounding box in an image.
:param rgbImg: RGB image to process. Shape: (height, width, 3)
:type rgbImg: numpy.ndarray
:param skipMulti: Skip image if more than one face detected.
:type skipMulti: bool
:return: The largest face bounding box in an image, or None.
:rtype: dlib.rectangle
"""
assert rgbImg is not None
faces = self.getAllFaceBoundingBoxes(rgbImg)
if len(faces) > 0:
if (not skipMulti and len(faces) > 0) or len(faces) == 1:
return max(faces, key=lambda rect: rect.width() * rect.height())
else:
return None
@ -143,7 +145,8 @@ class AlignDlib:
return list(map(lambda p: (p.x, p.y), points.parts()))
def align(self, imgDim, rgbImg, bb=None,
landmarks=None, landmarkIndices=INNER_EYES_AND_BOTTOM_LIP):
landmarks=None, landmarkIndices=INNER_EYES_AND_BOTTOM_LIP,
skipMulti=False):
r"""align(imgDim, rgbImg, bb=None, landmarks=None, landmarkIndices=INNER_EYES_AND_BOTTOM_LIP)
Transform and align a face in an image.
@ -160,6 +163,8 @@ class AlignDlib:
:type landmarks: list of (x,y) tuples
:param landmarkIndices: The indices to transform to.
:type landmarkIndices: list of ints
:param skipMulti: Skip image if more than one face detected.
:type skipMulti: bool
:return: The aligned RGB image. Shape: (imgDim, imgDim, 3)
:rtype: numpy.ndarray
"""
@ -168,7 +173,7 @@ class AlignDlib:
assert landmarkIndices is not None
if bb is None:
bb = self.getLargestFaceBoundingBox(rgbImg)
bb = self.getLargestFaceBoundingBox(rgbImg, skipMulti)
if bb is None:
return

View File

@ -112,7 +112,8 @@ def alignMain(args):
outRgb = None
else:
outRgb = align.align(args.size, rgb,
landmarkIndices=landmarkIndices)
landmarkIndices=landmarkIndices,
skipMulti=args.skipMulti)
if outRgb is None and args.verbose:
print(" + Unable to align.")
@ -157,6 +158,7 @@ if __name__ == '__main__':
default=96)
alignmentParser.add_argument('--fallbackLfw', type=str,
help="If alignment doesn't work, fallback to copying the deep funneled version from this directory..")
alignmentParser.add_argument('--skipMulti', action='store_true', help="Skip images with more than one face.")
alignmentParser.add_argument('--verbose', action='store_true')
args = parser.parse_args()

View File

@ -48,7 +48,7 @@ def main():
parser.add_argument('--imgDim', type=int,
help="Default image dimension.", default=96)
parser.add_argument('--cuda', action='store_true')
parser.add_argument('--threshold', type=int, default=0.9)
parser.add_argument('--threshold', type=float, default=0.9)
parser.add_argument('--delete', action='store_true', help='Delete the outliers.')
parser.add_argument('directory')