crop_face_in_video_for_FAS.py

import os
import cv2
import numpy as np
from tqdm import tqdm
import argparse

from facenet_pytorch import MTCNN

mtcnn = MTCNN(select_largest=True, min_face_size=64, post_process=False, device='cuda:0')

def get_file_paths(folder_path):
    file_paths = []
    for root, dirs, files in os.walk(folder_path):
        for file in files:
            file_paths.append(os.path.join(root, file))
    return file_paths

def crop_videos(in_folder, out_folder):
    print('Videos are in :', in_folder)
    os.makedirs(out_folder, exist_ok=True)
    skipped_imgs = []
    no_used_images = []

    video_names = get_file_paths(in_folder)
    for video_name in tqdm(video_names):
        video_path = os.path.join(in_folder, video_name)
        video_capture = cv2.VideoCapture(video_path)
        frame_count = 0
        while True:
            ret, frame = video_capture.read()
            if not ret:
                break
            frame_count += 1
            if frame_count % 10 == 0:  # 每10帧进行一次裁剪
                try:
                    # 使用MTCNN或其他人脸检测算法进行人脸检测
                    # 假设这里有一个名为mtcnn的人脸检测器
                    boxes, probs = mtcnn.detect(frame)
                except:
                    skipped_imgs.append((video_name, frame_count))
                    continue

                if boxes is None:
                    skipped_imgs.append((video_name, frame_count))
                    continue

                for i, box in enumerate(boxes):
                    x1, y1, x2, y2 = box
                    x1, y1 = max(int(x1), 0), max(int(y1), 0)
                    crop_img = frame[y1:int(y2), x1:int(x2)]

                    try:
                        cv2.imwrite(os.path.join(out_folder, f"{video_name.split('/')[-1]}_{frame_count}_{i}.png"), crop_img)
                    except:
                        skipped_imgs.append((video_name, frame_count))
                        pass

        video_capture.release()

    print(skipped_imgs)
    print(f"Images with no Face: {len(skipped_imgs)}")

def main():
    """
    Usage:
    python crop_face_in_video.py \
    --in_folder /home/Yanwei_Liu/New_Datasets/3600g/XXXXXXXXX/Videos_picture/val \
    --out_folder /home/Yanwei_Liu/New_Datasets/3600g/XXXXXXXXX/Videos_picture/val_crop
    """
    parser = argparse.ArgumentParser(description='Face detection and crop')
    parser.add_argument('--in_folder', type=str, default="/home/Yanwei_Liu/New_Datasets/3600g/XXXXXXXXX/content/train/real_video", help='folder with images')
    parser.add_argument('--out_folder', type=str, default="/home/Yanwei_Liu/New_Datasets/3600g/XXXXXXXXX/content/train/real_video_crop", help="folder to save aligned images")

    args = parser.parse_args()
    crop_videos(args.in_folder, args.out_folder)

if __name__ == "__main__":
    main()