-
Notifications
You must be signed in to change notification settings - Fork 10
/
demo.py
181 lines (147 loc) · 6.26 KB
/
demo.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
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
import face_recognition
import numpy as np
import cv2
import glob
import random
import string
import os
import math
import argparse
os.system('cls' if os.name=='nt' else 'clear')
parser = argparse.ArgumentParser();
parser.add_argument('-i', type=str, help='Image of target face to scan for.', required=True)
parser.add_argument('-v', type=str, help='Video to process', required=True)
parser.add_argument('-t', type=float, help='Tolerance of face detection, lower is stricter. (0.1-1.0)', default=0.6)
parser.add_argument('-f', type=int, help='Amount of frames per second to extract.', default=25)
parser.add_argument('-n', type=int, help='Number of frames with target face to save from each vid.', default=1000)
parser.add_argument('-s', type=int, help='Minimum KB size of images to keep in the faceset.', default=32)
args = vars(parser.parse_args())
if args['t'] > 1.0:
args['t'] = 1.0
elif args['t'] < 0.1:
args['t'] = 0.1
min_KB = args['s']
tol = args['t']
xfps = args['f']
targfname = args['i']
vid_dir = args['v']
faces_from_each_video = args['n']
if faces_from_each_video < 1:
faces_from_each_video = 1000
if min_KB < 1:
min_KB = 32
print("Target filename: " + targfname + ".")
print("Video input directory: " + vid_dir + ".")
print("Tolerance: " + str(tol) + ".")
print("Number of confirmed faces saved from each video: " + str(faces_from_each_video) + ".")
if(cv2.ocl.haveOpenCL()):
cv2.ocl.setUseOpenCL(True)
print("Using OpenCL: " + str(cv2.ocl.useOpenCL()) + ".")
target_image = face_recognition.load_image_file(targfname)
outdir = str(str(os.path.splitext(targfname)[0]) + "_output");
scanned_vids = str(str(os.path.splitext(targfname)[0]) + "_scanned_vids");
too_small = str(str(os.path.splitext(targfname)[0]) + "_too_small");
#check if output directories already exists, and if not, create it
os.makedirs(outdir, exist_ok=True)
os.makedirs(scanned_vids, exist_ok=True)
os.makedirs(too_small, exist_ok=True)
print("Output directory: " + outdir + ".")
print("Scanned videos will be moved to: " + scanned_vids + ".")
try:
target_encoding = face_recognition.face_encodings(target_image)[0]
except IndexError:
print("No face found in target image.")
raise SystemExit(0)
def random_string(length):
return ''.join(random.choice(string.ascii_letters) for m in range(length))
for zz in range(200):
try:
vid = random.choice(glob.glob(vid_dir + '/*.mp4'))
print("Now looking at video: " + vid)
input_video = cv2.VideoCapture(vid)
framenum = 0
vidheight = input_video.get(4)
vidwidth = input_video.get(3)
vidfps = input_video.get(cv2.CAP_PROP_FPS)
totalframes = input_video.get(cv2.CAP_PROP_FRAME_COUNT)
outputsize = 256, 256
if xfps > vidfps:
xfps = vidfps
print("Frame Width: " + str(vidwidth) + ", Height: " + str(vidheight) + ".")
known_faces = [
target_encoding
]
#switch to output directory
os.chdir(str(os.path.splitext(targfname)[0]) + "_output")
written = 1
while(input_video.isOpened()):
input_video.set(1, (framenum + (vidfps/xfps)))
framenum += vidfps/xfps
ret, frame = input_video.read()
if not ret:
break
percentage = (framenum/totalframes)*100
print("Checking frame " + str(int(framenum)) + "/" + str(int(totalframes)) + str(" (%.2f%%)" % percentage))
rgb_frame = frame[:, :, ::-1]
face_locations = face_recognition.face_locations(rgb_frame)
face_encodings = face_recognition.face_encodings(rgb_frame, face_locations)
for fenc, floc in zip(face_encodings, face_locations):
istarget = face_recognition.compare_faces(known_faces, fenc, tolerance=float(tol))
#if the face found matches the target
if istarget[0]:
top, right, bottom, left = floc
facefound = True
#squaring it up
if (bottom - top) > (right - left):
right = left + (bottom - top)
elif (right - left) > (bottom - top):
bottom = top + (right - left)
#calculating the diagonal of the cropped face for rotation purposes
#diagonal = math.sqrt(2*(bottom - top))
#padding = diagonal / 2
#alignment script causes images cropped "too closely" to get a bit fucky, so crop them less severely.
padding = (bottom - top)/2
if((top - padding >= 0) and (bottom + padding <= vidheight) and (left - padding >= 0) and (right + padding <= vidwidth)):
croppedframe = frame[int(top - padding):int(bottom + padding), int(left - padding):int(right + padding)]
#if the image is too small, resize it to outputsize
cheight, cwidth, cchannels = croppedframe.shape
if (cheight < 256) or (cwidth < 256):
croppedframe = cv2.resize(croppedframe, outputsize, interpolation=cv2.INTER_CUBIC)
print('Writing image ' + str(written) + '.')
cv2.imwrite(("vid_" + str(zz) + '_pic' + str(written) + '_' + random_string(15) + ".jpg"), croppedframe, [int(cv2.IMWRITE_JPEG_QUALITY), 98])
written += 1
if percentage > 99.9:
os.rename(vid, scanned_vids + '/vid' + str(zz) + '_' + random_string(5) + '.mp4')
break
if written > faces_from_each_video:
os.rename(vid, scanned_vids + '/vid' + str(zz) + '_' + random_string(5) + '.mp4')
break
input_video.release()
except ValueError:
print ("Scanning videos complete.")
pass
except IndexError:
pass
#Removes images under 32KB
counter = 0
low_quat = min_KB * 1000
for xx in (os.listdir(os.getcwd())):
if(os.path.getsize(xx)) < low_quat:
os.rename(xx, too_small + "/too small-" + str(counter) + random_string(15) + ".jpg")
print ("Moving " + str(xx) + " to the too small folder")
counter += 1
#Remove images with more than one face
print ("Now double checking there is only one face in each photo")
for yy in (os.listdir(os.getcwd())):
# Load the jpg file into a numpy array
image = face_recognition.load_image_file(yy)
# Find all the faces in the image using a pre-trained convolutional neural network.
# This method is more accurate than the default HOG model, but it's slower
# unless you have an nvidia GPU and dlib compiled with CUDA extensions. But if you do,
# this will use GPU acceleration and perform well.
# See also: find_faces_in_picture.py
face_locations = face_recognition.face_locations(image, number_of_times_to_upsample=0, model="cnn")
print("I found {} face(s) in this photograph.".format(len(face_locations)))
if not (len(face_locations)) == 1:
os.remove(yy)
print (str(yy) + ' was removed')