bz2 file support

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):

masks/masks.cfg

@@ -1,23 +1,23 @@
-[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
 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
 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

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)