def swap(self, image_fd, dataset, threshold=2.0, max_size=(1024,1024)):
target = faceswap.load_image(image_fd, max_size=max_size)
face_info = faceswap.extract_face_infos(target)
for info in face_info:
score, db_image_data = self.find(info.imghash, dataset)
if score < threshold:
db_image = faceswap.load_image(open(db_image_data['image_path']))
target = transform.faceswap(db_image, db_image_data, target, info)
return target
def _load_data(self, data_dir):
data = defaultdict(dict)
for dirpath, _, files in os.walk('./database/'):
db = os.path.basename(dirpath)
if db == 'cache':
cache_db = dirpath.split(os.path.sep)[-2]
print "Loading cache: ", cache_db
data[cache_db].update(self._load_cache(dirpath, files))
elif files:
print "Looking at DB: ", db
for filename in files:
if any(filename.endswith(f) for f in ('jpg', 'jpeg', 'png')):
fullpath = os.path.join(dirpath, filename)
image = faceswap.load_image(open(fullpath))
basename = os.path.splitext(filename)[0]
cache_loc = os.path.join(dirpath, "cache", basename + '_0.pkl')
if not os.path.exists(cache_loc):
print "\tAnalysing: ", filename
extracted_data = self._parse_image(image, basename, db, fullpath)
data[db].update(extracted_data)
for db, raw_data in data.iteritems():
if raw_data:
print "Database {} has {} items".format(db, len(raw_data))
db_items = raw_data.values()
self.image_lookup[db] = db_items
self.hashes_kdtree[db] = cKDTree([data['imghash'] for data in db_items])
