Comparison demo: Use all combinations of images.

This commit is contained in:
Brandon Amos 2015-10-03 19:23:13 -04:00
parent 28891143a1
commit 402b453f05
1 changed files with 25 additions and 19 deletions

View File

@ -23,6 +23,7 @@ import time
start = time.time() start = time.time()
import argparse import argparse
import cv2 import cv2
import itertools
import os import os
import numpy as np import numpy as np
@ -42,8 +43,7 @@ facenetModelDir = os.path.join(modelDir, 'facenet')
parser = argparse.ArgumentParser() parser = argparse.ArgumentParser()
parser.add_argument('img1', type=str, help="Input image 1.") parser.add_argument('imgs', type=str, nargs='+', help="Input images.")
parser.add_argument('img2', type=str, help="Input image 2.")
parser.add_argument('--dlibFaceMean', type=str, help="Path to dlib's face predictor.", parser.add_argument('--dlibFaceMean', type=str, help="Path to dlib's face predictor.",
default=os.path.join(dlibModelDir, "mean.csv")) default=os.path.join(dlibModelDir, "mean.csv"))
parser.add_argument('--dlibFacePredictor', type=str, help="Path to dlib's face predictor.", parser.add_argument('--dlibFacePredictor', type=str, help="Path to dlib's face predictor.",
@ -55,6 +55,7 @@ parser.add_argument('--networkModel', type=str, help="Path to Torch network mode
default=os.path.join(facenetModelDir, 'nn4.v1.t7')) default=os.path.join(facenetModelDir, 'nn4.v1.t7'))
parser.add_argument('--imgDim', type=int, help="Default image dimension.", default=96) parser.add_argument('--imgDim', type=int, help="Default image dimension.", default=96)
parser.add_argument('--cuda', type=bool, default=False) parser.add_argument('--cuda', type=bool, default=False)
parser.add_argument('--verbose', type=bool, default=False)
args = parser.parse_args() args = parser.parse_args()
@ -62,43 +63,48 @@ sys.path.append(args.dlibRoot)
import dlib import dlib
from facenet.alignment import NaiveDlib # Depends on dlib. from facenet.alignment import NaiveDlib # Depends on dlib.
if args.verbose:
print("Argument parsing and loading libraries took {} seconds.".format(time.time()-start)) print("Argument parsing and loading libraries took {} seconds.".format(time.time()-start))
start = time.time() start = time.time()
align = NaiveDlib(args.dlibFaceMean, args.dlibFacePredictor) align = NaiveDlib(args.dlibFaceMean, args.dlibFacePredictor)
net = facenet.TorchWrap(args.networkModel, imgDim=args.imgDim, cuda=args.cuda) net = facenet.TorchWrap(args.networkModel, imgDim=args.imgDim, cuda=args.cuda)
if args.verbose:
print("Loading the dlib and FaceNet models took {} seconds.".format(time.time()-start)) print("Loading the dlib and FaceNet models took {} seconds.".format(time.time()-start))
def getRep(imgPath): def getRep(imgPath):
global i if args.verbose:
print("Processing {}.".format(imgPath)) print("Processing {}.".format(imgPath))
img = cv2.imread(imgPath) img = cv2.imread(imgPath)
if img is None: if img is None:
raise Exception("Unable to load image: {}".format(imgPath)) raise Exception("Unable to load image: {}".format(imgPath))
if args.verbose:
print(" + Original size: {}".format(img.shape)) print(" + Original size: {}".format(img.shape))
start = time.time() start = time.time()
bb = align.getLargestFaceBoundingBox(img) bb = align.getLargestFaceBoundingBox(img)
if bb is None: if bb is None:
raise Exception("Unable to find a face: {}".format(imgPath)) raise Exception("Unable to find a face: {}".format(imgPath))
if args.verbose:
print(" + Face detection took {} seconds.".format(time.time()-start)) print(" + Face detection took {} seconds.".format(time.time()-start))
start = time.time() start = time.time()
alignedFace = align.alignImg("affine", args.imgDim, img, bb) alignedFace = align.alignImg("affine", args.imgDim, img, bb)
if alignedFace is None: if alignedFace is None:
raise Exception("Unable to align image: {}".format(imgPath)) raise Exception("Unable to align image: {}".format(imgPath))
if args.verbose:
print(" + Face alignment took {} seconds.".format(time.time()-start)) print(" + Face alignment took {} seconds.".format(time.time()-start))
start = time.time() start = time.time()
t = '/tmp/facenet-compare.png' rep = net.forwardImage(alignedFace)
cv2.imwrite(t, alignedFace) if args.verbose:
rep = np.array(net.forward(t))
os.remove(t)
print(" + FaceNet forward pass took {} seconds.".format(time.time()-start)) print(" + FaceNet forward pass took {} seconds.".format(time.time()-start))
print("Representation:") print("Representation:")
print(rep) print(rep)
print("-----\n") print("-----\n")
return rep return rep
d = getRep(args.img1) - getRep(args.img2) for (img1, img2) in itertools.combinations(args.imgs, 2):
print("Squared l2 distance between representations: {}".format(np.dot(d, d))) d = getRep(img1) - getRep(img2)
print("Comparing {} with {}.".format(img1, img2))
print(" + Squared l2 distance between representations: {}".format(np.dot(d, d)))