Spaces:
Runtime error
Runtime error
import time | |
import imutils | |
from imutils import paths | |
import face_recognition | |
import pickle | |
import cv2 | |
import os | |
import argparse | |
def augment_image(file_path): | |
""" | |
Augment the image by rotating it by 10, -10, 30, -30 degrees, flipping it horizontally and increasing/decreasing the brightness | |
params: | |
file_path: path to the image file | |
returns: | |
list of 8 augmented images | |
""" | |
image = cv2.imread(file_path) | |
images = [] | |
# convert the input image from BGR to RGB | |
rgb_image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) | |
# add the original image to the list of images | |
images.append(rgb_image) | |
# rotate the image by 10 degrees and add it to the list of images | |
images.append(imutils.rotate_bound(rgb_image, 10)) | |
# rotate the image by -10 degrees and add it to the list of images | |
images.append(imutils.rotate_bound(rgb_image, -10)) | |
# rotate the image by 30 degrees and add it to the list of images | |
images.append(imutils.rotate_bound(rgb_image, 30)) | |
# rotate the image by -30 degrees and add it to the list of images | |
images.append(imutils.rotate_bound(rgb_image, -30)) | |
# flip the image horizontally and add it to the list of images | |
images.append(cv2.flip(rgb_image, 1)) | |
#increase the brightness of the image and add it to the list of images | |
images.append(cv2.convertScaleAbs(rgb_image, alpha=1.5, beta=0)) | |
#decrease the brightness of the image and add it to the list of images | |
images.append(cv2.convertScaleAbs(rgb_image, alpha=0.5, beta=0)) | |
return images | |
def preprocess(): | |
dataset = './dataset' # path to input dataset directory | |
encodings_file = './encodings.pkl' #path to output encodings file | |
imagePaths = list(paths.list_images(dataset)) | |
known_encodings = [] | |
known_names = [] | |
s = time.time() | |
print(imagePaths) | |
for i, imagePath in enumerate(imagePaths): | |
images = augment_image(imagePath) | |
name = imagePath.split(os.path.sep)[-2] | |
print(f"processing image [{name}] {i+1}/{len(imagePaths)}") | |
for l in images: | |
boxes = face_recognition.face_locations(l, model='hog') | |
encodings = face_recognition.face_encodings(l, boxes) | |
# creating the training set | |
for encoding in encodings: | |
known_encodings.append(encoding) | |
known_names.append(name) | |
e = time.time() | |
print(f"Encoding dataset took: {(e-s)} seconds") | |
data = {"encodings": known_encodings, "names": known_names} | |
f = open(encodings_file, "wb") | |
f.write(pickle.dumps(data)) | |
f.close() | |
return encodings_file | |
# replace the above few lines with the script to upload this to hopsworks | |
if __name__ == '__main__': | |
preprocess() |