Skip to content

Instantly share code, notes, and snippets.

@e96031413
Created March 11, 2024 02:21
Show Gist options
  • Save e96031413/74e1219ccf13db7ccfbc8fb481923b80 to your computer and use it in GitHub Desktop.
Save e96031413/74e1219ccf13db7ccfbc8fb481923b80 to your computer and use it in GitHub Desktop.
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()
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment