diff --git a/dlib_models/readme.txt b/dlib_models/readme.txt deleted file mode 100644 index d233962..0000000 --- a/dlib_models/readme.txt +++ /dev/null @@ -1 +0,0 @@ -Save the dlib .dat model here \ No newline at end of file diff --git a/mask_the_face.py b/mask_the_face.py index 51f5693..0caf9ef 100644 --- a/mask_the_face.py +++ b/mask_the_face.py @@ -5,14 +5,14 @@ import argparse import dlib from utils.aux_functions import * -import random #TODO: -# 1. surgical_green, surgical_blue --> one surgical +# 1. Done: surgical_green, surgical_blue --> one surgical # 2. left mask and right mask --> one angled mask # 3. Done: MFR2 Dataset script # 4. Done: Organize and Upload MFR2 dataset # 5. Done: Dlib based detector +# 6. Done: Download dlib model # Command-line input setup parser = argparse.ArgumentParser( @@ -27,7 +27,7 @@ parser.add_argument( "--mask_type", type=str, - default="N95", + default="surgical", choices=['surgical', 'N95', 'KN95', 'cloth', 'gas', 'inpaint', 'random', 'all'], help="Type of the mask to be applied. Available options: all, surgical_blue, surgical_green, N95, cloth", ) @@ -49,7 +49,7 @@ parser.add_argument( "--color", type=str, - default="", + default="#0473e2", help="Hex color value that need to be overlayed to the mask", ) @@ -63,7 +63,8 @@ parser.add_argument( "--code", type = str, - default="cloth-masks/textures/check/check_4.jpg, cloth-#e54294, cloth-#ff0000, cloth, cloth-masks/textures/others/heart_1.png, cloth-masks/textures/fruits/pineapple.png, N95, surgical_blue, surgical_green", + # default="cloth-masks/textures/check/check_4.jpg, cloth-#e54294, cloth-#ff0000, cloth, cloth-masks/textures/others/heart_1.png, cloth-masks/textures/fruits/pineapple.png, N95, surgical_blue, surgical_green", + default="", help="Generate specific formats", ) @@ -84,7 +85,11 @@ # Set up dlib face detector and predictor args.detector = dlib.get_frontal_face_detector() -args.predictor = dlib.shape_predictor("dlib_models/shape_predictor_68_face_landmarks.dat") +path_to_dlib_model = 'dlib_models/shape_predictor_68_face_landmarks.dat' +if not os.path.exists(path_to_dlib_model): + download_dlib_model() + +args.predictor = dlib.shape_predictor(path_to_dlib_model) # Extract data from code mask_code = "".join(args.code.split()).split(',') @@ -116,7 +121,8 @@ path, dirs, files = os.walk(args.path).__next__() file_count = len(files) dirs_count = len(dirs) - print_orderly("Masking image files", 57) + if len(files)>0: + print_orderly("Masking image files", 60) # Process files in the directory if any for f in tqdm(files): @@ -149,7 +155,7 @@ img = masked_image[i] cv2.imwrite(w_path, img) - print_orderly("Masking image directories", 57) + print_orderly("Masking image directories", 60) # Process directories withing the path provided for d in tqdm(dirs): diff --git a/masks/masks.cfg b/masks/masks.cfg index edb8f96..92e9be0 100644 --- a/masks/masks.cfg +++ b/masks/masks.cfg @@ -1,5 +1,5 @@ -[surgical_blue] -template: masks/templates/surgical_blue.png +[surgical] +template: masks/templates/surgical.png mask_a: 21, 97 mask_b: 307, 22 mask_c: 600, 99 @@ -7,8 +7,8 @@ mask_d: 25, 322 mask_e: 295, 470 mask_f: 600, 323 -[surgical_blue_left] -template: masks/templates/surgical_blue_left.png +[surgical_left] +template: masks/templates/surgical_left.png mask_a: 39, 27 mask_b: 130, 9 mask_c: 567, 20 @@ -16,8 +16,8 @@ mask_d: 87, 207 mask_e: 168, 302 mask_f: 568, 202 -[surgical_blue_right] -template: masks/templates/surgical_blue_right.png +[surgical_right] +template: masks/templates/surgical_right.png mask_a: 3, 20 mask_b: 440, 9 mask_c: 531, 27 diff --git a/utils/aux_functions.py b/utils/aux_functions.py index 1196bb9..4b5ed35 100644 --- a/utils/aux_functions.py +++ b/utils/aux_functions.py @@ -11,6 +11,31 @@ import random from utils.create_mask import texture_the_mask, color_the_mask from imutils import face_utils +import requests +from zipfile import ZipFile +from tqdm import tqdm +import bz2, shutil + +def download_dlib_model(): + print_orderly('Get dlib model', 60) + dlib_model_link = "http://dlib.net/files/shape_predictor_68_face_landmarks.dat.bz2" + print('Downloading dlib model...') + r = requests.get(dlib_model_link, stream=True) + print("Zip file size: ", np.round(len(r.content)/1024/1024, 2), "MB") + destination = 'dlib_models'+os.path.sep+'shape_predictor_68_face_landmarks.dat.bz2' + if not os.path.exists(destination.rsplit(os.path.sep, 1)[0]): + os.mkdir(destination.rsplit(os.path.sep, 1)[0]) + print('Saving dlib model...') + with open(destination, 'wb') as fd: + for chunk in r.iter_content(chunk_size=32678): + fd.write(chunk) + print('Extracting dlib model...') + with bz2.BZ2File(destination) as fr, open("dlib_models/shape_predictor_68_face_landmarks.dat", "wb") as fw: + shutil.copyfileobj(fr, fw) + print('Saved: ', destination) + print_orderly('done', 60) + + os.remove(destination) def get_line(face_landmark, image, type="eye", debug=False): pil_image = Image.fromarray(image) @@ -630,7 +655,7 @@ def get_available_mask_types(config_filename="masks/masks.cfg"): def print_orderly(str, n): - print("") + # print("") hyphens = "-" * int((n - len(str)) / 2) str_p = hyphens + " " + str + " " + hyphens hyphens_bar = "-" * len(str_p)