Added ability to skip images with more than one face; fixed threshold type
This commit is contained in:
parent
5b46a5fcb4
commit
d18017a195
|
@ -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
|
||||
|
||||
|
|
|
@ -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()
|
||||
|
|
|
@ -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')
|
||||
|
||||
|
|
Loading…
Reference in New Issue