Add checks for 'None' to resolve #25.

This commit is contained in:
Brandon Amos 2016-01-11 17:24:52 -05:00
parent bddc2538ed
commit b84d57a6df
4 changed files with 40 additions and 2 deletions

View File

@ -85,6 +85,9 @@ class AlignDlib:
:param facePredictor: The path to dlib's
:type facePredictor: str
"""
assert facePredictor is not None
self.detector = dlib.get_frontal_face_detector()
self.predictor = dlib.shape_predictor(facePredictor)
@ -97,6 +100,9 @@ class AlignDlib:
:return: All face bounding boxes in an image.
:rtype: dlib.rectangles
"""
assert rgbImg is not None
try:
return self.detector(rgbImg, 1)
except Exception as e:
@ -113,6 +119,9 @@ class AlignDlib:
: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:
return max(faces, key=lambda rect: rect.width() * rect.height())
@ -130,6 +139,10 @@ class AlignDlib:
:return: Detected landmark locations.
:rtype: list of (x,y) tuples
"""
assert rgbImg is not None
assert bb is not None
points = self.predictor(rgbImg, bb)
return list(map(lambda p: (p.x, p.y), points.parts()))
@ -154,6 +167,11 @@ class AlignDlib:
:return: The aligned RGB image. Shape: (imgDim, imgDim, 3)
:rtype: numpy.ndarray
"""
assert imgDim is not None
assert rgbImg is not None
assert landmarkIndices is not None
if bb is None:
bb = self.getLargestFaceBoundingBox(rgbImg)
if bb is None:

View File

@ -35,6 +35,11 @@ class Image:
:param path: Path to the image on disk.
:type path: str
"""
assert cls is not None
assert name is not None
assert path is not None
self.cls = cls
self.name = name
self.path = path
@ -98,6 +103,9 @@ def iterImgs(directory):
:type directory: str
:return: An iterator over Image objects.
"""
assert directory is not None
exts = [".jpg", ".png"]
for subdir, dirs, files in os.walk(directory):

View File

@ -13,6 +13,9 @@ def mkdirP(path):
:param path: The directory to create.
:type path: str
"""
assert path is not None
try:
os.makedirs(path)
except OSError as exc: # Python >2.5

View File

@ -53,6 +53,11 @@ class TorchNeuralNet:
:param cuda: Flag to use CUDA in the subprocess.
:type cuda: bool
"""
assert model is not None
assert imgDim is not None
assert cuda is not None
self.cmd = ['/usr/bin/env', 'th', os.path.join(myDir, 'openface_server.lua'),
'-model', model, '-imgDim', str(imgDim)]
if cuda:
@ -74,6 +79,9 @@ class TorchNeuralNet:
:return: Vector of features extracted with the neural network.
:rtype: numpy.ndarray
"""
assert imgPath is not None
rc = self.p.poll()
if rc is not None and rc != 0:
raise Exception("""
@ -134,8 +142,9 @@ stderr: {}
:return: Vector of features extracted from the neural network.
:rtype: numpy.ndarray
"""
if rgbImg is None:
raise Exception("rgbImg=None passed into forward")
assert rgbImg is not None
t = '/tmp/openface-torchwrap-{}.png'.format(
binascii.b2a_hex(os.urandom(8)))
bgrImg = cv2.cvtColor(rgbImg, cv2.COLOR_RGB2BGR)