-
Notifications
You must be signed in to change notification settings - Fork 0
/
filter_faces.py
129 lines (110 loc) · 5.04 KB
/
filter_faces.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
from argparse import ArgumentParser
import insightface
import cv2
import os
import logging
import traceback
import shutil
from tqdm import tqdm, tqdm_gui
import time
import numpy as np
import multiprocessing
from multiprocessing import Pool, cpu_count
from p_tqdm import p_map
import sys
import utils
def run(paths):
src, out = paths
if not os.path.exists(out):
img = cv2.imread(src)
img = cv2.resize(img, (255, 255))
bbox, _ = model.detect(img, threshold=0.5, scale=1.0)
if len(bbox) > 0:
try:
if CROP_FACES:
if DUPLICATE:
if MOVE_IMGS:
shutil.move(src, out)
else:
utils.copyFile(src, out)
for ith_face, box in enumerate(bbox):
x,y,w,h,_ = list(map(int, box))
imgCrop = img[y:y+h,x:x+w]
if not imgCrop.empty():
out += f"face_{str(ith_face)}.jpg"
cv2.imwrite(out, imgCrop)
if MOVE_IMGS:
shutil.move(src, out)
else:
utils.copyFile(src, out)
# utils.write("Filtered images: {} | Saved images percentage: {}".format(num_imgs_filtered, saved_imgs_percentage))
except Exception as e:
utils.write(str(e))
else:
utils.write("ALREADY_EXISTS")
def main(args):
model = insightface.model_zoo.get_model('retinaface_r50_v1')
model.prepare(ctx_id = -1, nms=0.4)
pathsGenerator = utils.yieldPaths(args["INPUT_BASE_DIR"], args["OUTPUT_PATH"], flat = FLAT)
if not FLAT:
utils.copyDirectoryStructure(args["INPUT_BASE_DIR"], args["OUTPUT_PATH"])
logging.info("Starting...")
with Pool(processes=NUM_CPUS) as p:
with tqdm(total=NUM_FILES,
desc=" ",
unit="Images") as pbar:
for i, _ in enumerate(p.imap_unordered(run, pathsGenerator)):
pbar.update()
percentage = round((i/NUM_FILES)*100, 3)
utils.write(f"Images processed:{i} | Percentage: {percentage}%")
if __name__ == '__main__':
# Initialize parser
parser = ArgumentParser(
description="Script for detecting faces in a given folder and its subdirectories"
)
parser.add_argument("-in", "--input-path",
type=str,
required=True,
dest="INPUT_BASE_DIR",
help="Path to the directory where images or folders of images are\n")
parser.add_argument("-out","--output-path",
type=str,
required=True,
dest = "OUTPUT_PATH",
help="Path of the folder where faces images will be saved\n")
parser.add_argument("-move", "--move-kept-images",
action="store_true",
default=False,
dest = "move_images",
help = "Whether to move kept images from [-in, --input-path] to [-out, --output-path] in such a way that in the remaining images in [-in --input-path] are the ones that did not apply the criteria.")
parser.add_argument("-crop","--crop-faces",
action='store_true',
dest="crop_faces",
default=False,
help="Crop faces detected in images and save each one\n")
parser.add_argument("-flat", "--same-out-dir",
action='store_true',
dest="save_in_same_output_folder",
default=False,
help="Whether to save all images in dirctory specified in -out --output-path and not imitate directory structure from the path specified in -indir --input-base-dir\n")
parser.add_argument("-duplicate", "--duplicate-img-faces",
action="store_true",
dest="duplicate_img_of_faces",
default=False,
help="Whether to save the original images of the extracted faces also. Only valid if -crop --crop-faces is passed as argument")
logging.basicConfig(level=logging.INFO)
logging.info(" Preparing model...")
model = insightface.model_zoo.get_model('retinaface_r50_v1')
model.prepare(ctx_id = -1, nms=0.4)
args = vars(parser.parse_args())
DUPLICATE = args["duplicate_img_of_faces"]
MOVE_IMGS = args["move_images"]
FLAT = args["save_in_same_output_folder"]
CROP_FACES= args["crop_faces"]
NUM_FILES = utils.countFiles(args["INPUT_BASE_DIR"])
NUM_CPUS = cpu_count()
num_imgs_filtered = 0
processed_imgs = 0
logging.info(f"{str(NUM_FILES)} files found")
logging.info(f"Working with {NUM_CPUS} cpus")
main(args)