| import os |
| import cv2 |
| import time |
| import glob |
| import argparse |
| import face_alignment |
| import numpy as np |
| from PIL import Image |
| from tqdm import tqdm |
| from itertools import cycle |
|
|
| from torch.multiprocessing import Pool, Process, set_start_method |
|
|
| class KeypointExtractor(): |
| def __init__(self, device): |
| self.detector = face_alignment.FaceAlignment(face_alignment.LandmarksType._2D, |
| device=device) |
|
|
| def extract_keypoint(self, images, name=None, info=True): |
| if isinstance(images, list): |
| keypoints = [] |
| if info: |
| i_range = tqdm(images,desc='landmark Det:') |
| else: |
| i_range = images |
|
|
| for image in i_range: |
| current_kp = self.extract_keypoint(image) |
| if np.mean(current_kp) == -1 and keypoints: |
| keypoints.append(keypoints[-1]) |
| else: |
| keypoints.append(current_kp[None]) |
|
|
| keypoints = np.concatenate(keypoints, 0) |
| np.savetxt(os.path.splitext(name)[0]+'.txt', keypoints.reshape(-1)) |
| return keypoints |
| else: |
| while True: |
| try: |
| keypoints = self.detector.get_landmarks_from_image(np.array(images))[0] |
| break |
| except RuntimeError as e: |
| if str(e).startswith('CUDA'): |
| print("Warning: out of memory, sleep for 1s") |
| time.sleep(1) |
| else: |
| print(e) |
| break |
| except TypeError: |
| print('No face detected in this image') |
| shape = [68, 2] |
| keypoints = -1. * np.ones(shape) |
| break |
| if name is not None: |
| np.savetxt(os.path.splitext(name)[0]+'.txt', keypoints.reshape(-1)) |
| return keypoints |
|
|
| def read_video(filename): |
| frames = [] |
| cap = cv2.VideoCapture(filename) |
| while cap.isOpened(): |
| ret, frame = cap.read() |
| if ret: |
| frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) |
| frame = Image.fromarray(frame) |
| frames.append(frame) |
| else: |
| break |
| cap.release() |
| return frames |
|
|
| def run(data): |
| filename, opt, device = data |
| os.environ['CUDA_VISIBLE_DEVICES'] = device |
| kp_extractor = KeypointExtractor() |
| images = read_video(filename) |
| name = filename.split('/')[-2:] |
| os.makedirs(os.path.join(opt.output_dir, name[-2]), exist_ok=True) |
| kp_extractor.extract_keypoint( |
| images, |
| name=os.path.join(opt.output_dir, name[-2], name[-1]) |
| ) |
|
|
| if __name__ == '__main__': |
| set_start_method('spawn') |
| parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter) |
| parser.add_argument('--input_dir', type=str, help='the folder of the input files') |
| parser.add_argument('--output_dir', type=str, help='the folder of the output files') |
| parser.add_argument('--device_ids', type=str, default='0,1') |
| parser.add_argument('--workers', type=int, default=4) |
|
|
| opt = parser.parse_args() |
| filenames = list() |
| VIDEO_EXTENSIONS_LOWERCASE = {'mp4'} |
| VIDEO_EXTENSIONS = VIDEO_EXTENSIONS_LOWERCASE.union({f.upper() for f in VIDEO_EXTENSIONS_LOWERCASE}) |
| extensions = VIDEO_EXTENSIONS |
| |
| for ext in extensions: |
| os.listdir(f'{opt.input_dir}') |
| print(f'{opt.input_dir}/*.{ext}') |
| filenames = sorted(glob.glob(f'{opt.input_dir}/*.{ext}')) |
| print('Total number of videos:', len(filenames)) |
| pool = Pool(opt.workers) |
| args_list = cycle([opt]) |
| device_ids = opt.device_ids.split(",") |
| device_ids = cycle(device_ids) |
| for data in tqdm(pool.imap_unordered(run, zip(filenames, args_list, device_ids))): |
| None |
|
|
