Add files via upload

mask_classifier/Data_Generator/loop_through_folder.py

@@ -0,0 +1,18 @@
+import cv2
+import os
+from mask import create_mask
+
+
+folder_path = "/home/preeth/Downloads"
+#dist_path = "/home/preeth/Downloads"
+
+#c = 0
+images = [os.path.join(folder_path, f) for f in os.listdir(folder_path) if os.path.isfile(os.path.join(folder_path, f))]
+for i in range(len(images)):
+    print("the path of the image is", images[i])
+    #image = cv2.imread(images[i])
+    #c = c + 1
+    create_mask(images[i])
+
+
+

mask_classifier/Data_Generator/mask.py

@@ -0,0 +1,170 @@
+import os
+import sys
+import random
+import argparse
+import numpy as np
+from PIL import Image, ImageFile
+
+__version__ = '0.3.0'
+
+
+IMAGE_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'images')
+DEFAULT_IMAGE_PATH = os.path.join(IMAGE_DIR, 'default-mask.png')
+BLACK_IMAGE_PATH = os.path.join(IMAGE_DIR, 'black-mask.png')
+BLUE_IMAGE_PATH = os.path.join(IMAGE_DIR, 'blue-mask.png')
+RED_IMAGE_PATH = os.path.join(IMAGE_DIR, 'red-mask.png')
+
+
+def cli():
+    parser = argparse.ArgumentParser(description='Wear a face mask in the given picture.')
+    parser.add_argument('pic_path', help='Picture path.')
+    parser.add_argument('--show', action='store_true', help='Whether show picture with mask or not.')
+    parser.add_argument('--model', default='hog', choices=['hog', 'cnn'], help='Which face detection model to use.')
+    group = parser.add_mutually_exclusive_group()
+    group.add_argument('--black', action='store_true', help='Wear black mask')
+    group.add_argument('--blue', action='store_true', help='Wear blue mask')
+    group.add_argument('--red', action='store_true', help='Wear red mask')
+    args = parser.parse_args()
+
+    pic_path = args.pic_path
+    if not os.path.exists(args.pic_path):
+        print(f'Picture {pic_path} not exists.')
+        sys.exit(1)
+
+    if args.black:
+        mask_path = BLACK_IMAGE_PATH
+    elif args.blue:
+        mask_path = BLUE_IMAGE_PATH
+    elif args.red:
+        mask_path = RED_IMAGE_PATH
+    else:
+        mask_path = DEFAULT_IMAGE_PATH
+
+    FaceMasker(pic_path, mask_path, args.show, args.model).mask()
+
+
+def create_mask(image_path):
+    pic_path = image_path
+    mask_path = "/media/preeth/Data/prajna_files/mask_creator/face_mask/images/blue-mask.png"
+    show = False
+    model = "hog"
+    FaceMasker(pic_path, mask_path, show, model).mask()
+
+
+
+class FaceMasker:
+    KEY_FACIAL_FEATURES = ('nose_bridge', 'chin')
+
+    def __init__(self, face_path, mask_path, show=False, model='hog'):
+        self.face_path = face_path
+        self.mask_path = mask_path
+        self.show = show
+        self.model = model
+        self._face_img: ImageFile = None
+        self._mask_img: ImageFile = None
+
+    def mask(self):
+        import face_recognition
+
+        face_image_np = face_recognition.load_image_file(self.face_path)
+        face_locations = face_recognition.face_locations(face_image_np, model=self.model)
+        face_landmarks = face_recognition.face_landmarks(face_image_np, face_locations)
+        self._face_img = Image.fromarray(face_image_np)
+        self._mask_img = Image.open(self.mask_path)
+
+        found_face = False
+        for face_landmark in face_landmarks:
+            # check whether facial features meet requirement
+            skip = False
+            for facial_feature in self.KEY_FACIAL_FEATURES:
+                if facial_feature not in face_landmark:
+                    skip = True
+                    break
+            if skip:
+                continue
+
+            # mask face
+            found_face = True
+            self._mask_face(face_landmark)
+
+        if found_face:
+            if self.show:
+                self._face_img.show()
+
+            # save
+            self._save()
+        else:
+            print('Found no face.')
+
+    def _mask_face(self, face_landmark: dict):
+        nose_bridge = face_landmark['nose_bridge']
+        nose_point = nose_bridge[len(nose_bridge) * 1 // 4]
+        nose_v = np.array(nose_point)
+
+        chin = face_landmark['chin']
+        chin_len = len(chin)
+        chin_bottom_point = chin[chin_len // 2]
+        chin_bottom_v = np.array(chin_bottom_point)
+        chin_left_point = chin[chin_len // 8]
+        chin_right_point = chin[chin_len * 7 // 8]
+
+        # split mask and resize
+        width = self._mask_img.width
+        height = self._mask_img.height
+        width_ratio = 1.2
+        new_height = int(np.linalg.norm(nose_v - chin_bottom_v))
+
+        # left
+        mask_left_img = self._mask_img.crop((0, 0, width // 2, height))
+        mask_left_width = self.get_distance_from_point_to_line(chin_left_point, nose_point, chin_bottom_point)
+        mask_left_width = int(mask_left_width * width_ratio)
+        mask_left_img = mask_left_img.resize((mask_left_width, new_height))
+
+        # right
+        mask_right_img = self._mask_img.crop((width // 2, 0, width, height))
+        mask_right_width = self.get_distance_from_point_to_line(chin_right_point, nose_point, chin_bottom_point)
+        mask_right_width = int(mask_right_width * width_ratio)
+        mask_right_img = mask_right_img.resize((mask_right_width, new_height))
+
+        # merge mask
+        size = (mask_left_img.width + mask_right_img.width, new_height)
+        mask_img = Image.new('RGBA', size)
+        mask_img.paste(mask_left_img, (0, 0), mask_left_img)
+        mask_img.paste(mask_right_img, (mask_left_img.width, 0), mask_right_img)
+
+        # rotate mask
+        angle = np.arctan2(chin_bottom_point[1] - nose_point[1], chin_bottom_point[0] - nose_point[0])
+        rotated_mask_img = mask_img.rotate(angle, expand=True)
+
+        # calculate mask location
+        center_x = (nose_point[0] + chin_bottom_point[0]) // 2
+        center_y = (nose_point[1] + chin_bottom_point[1]) // 2
+
+        offset = mask_img.width // 2 - mask_left_img.width
+        radian = angle * np.pi / 180
+        box_x = center_x + int(offset * np.cos(radian)) - rotated_mask_img.width // 2
+        box_y = center_y + int(offset * np.sin(radian)) - rotated_mask_img.height // 2
+
+        # add mask
+        self._face_img.paste(mask_img, (box_x, box_y), mask_img)
+
+    def _save(self):
+        path_splits = os.path.splitext(self.face_path)
+        new_face_path = path_splits[0] + '-with-mask' + path_splits[1]
+        self._face_img.save(new_face_path)
+        print(f'Save to {new_face_path}')
+
+    @staticmethod
+    def get_distance_from_point_to_line(point, line_point1, line_point2):
+        distance = np.abs((line_point2[1] - line_point1[1]) * point[0] +
+                          (line_point1[0] - line_point2[0]) * point[1] +
+                          (line_point2[0] - line_point1[0]) * line_point1[1] +
+                          (line_point1[1] - line_point2[1]) * line_point1[0]) / \
+                   np.sqrt((line_point2[1] - line_point1[1]) * (line_point2[1] - line_point1[1]) +
+                           (line_point1[0] - line_point2[0]) * (line_point1[0] - line_point2[0]))
+        return int(distance)
+
+
+if __name__ == '__main__':
+    #cli()
+    create_mask(image_path)