def main():
torch.set_grad_enabled(False)
input_file = sys.argv[1]
output_dir = sys.argv[2]
reader = cv2.VideoCapture(input_file)
face_detector = FaceDetector()
face_detector.load_checkpoint("RetinaFace-Resnet50-fixed.pth")
for idx in itertools.count():
success, img = reader.read()
if not success:
break
boxes, landms = face_detector.detect(img)
if boxes.shape[0] == 0:
continue
areas = (boxes[:, 3] - boxes[:, 1]) * (boxes[:, 2] - boxes[:, 0])
max_face_idx = areas.argmax()
landm = landms[max_face_idx]
landmarks = landm.numpy().reshape(5, 2).astype(np.int)
img = norm_crop(img, landmarks, image_size=320)
aligned = Image.fromarray(img[:, :, ::-1])
out_path = os.path.join(output_dir, "%03d.jpg" % idx)
aligned.save(out_path)
def iter_one_face(self):
for fname in self.file_list:
path = os.path.join(self.video_dir, fname)
reader = cv2.VideoCapture(path)
face_count = 0
while True:
for _ in range(self.frame_skip):
reader.grab()
success, img = reader.read()
# 检测这一帧是否存在
if not success:
break
boxes, landms = self.face_detector.detect(img)
if boxes.shape[0] == 0:
continue
areas = (boxes[:, 3] - boxes[:, 1]) * (boxes[:, 2] -
boxes[:, 0])
order = areas.argmax()
boxes = boxes[order]
landms = landms[order]
# Crop faces
landmarks = landms.numpy().reshape(5, 2).astype(np.int)
img = norm_crop(img, landmarks, image_size=320)
aligned = Image.fromarray(img[:, :, ::-1])
if self.transform:
aligned = self.transform(aligned)
yield fname, aligned
# Early stop
face_count += 1
if face_count == self.face_limit:
break
reader.release()