File size: 9,488 Bytes
397bbeb
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
import os, sys
import numpy as np
from tqdm import tqdm, trange
import deep_3drecon
from moviepy.editor import VideoFileClip
from utils.commons.multiprocess_utils import multiprocess_run_tqdm, multiprocess_run
from utils.commons.meters import Timer
from decord import VideoReader
from decord import cpu, gpu
from utils.commons.face_alignment_utils import mediapipe_lm478_to_face_alignment_lm68
import mediapipe
import cv2

sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))

# fa = face_alignment.FaceAlignment(face_alignment.LandmarksType._2D, network_size=4, device='cuda')
mp_face_mesh = mediapipe.solutions.face_mesh
face_reconstructor = deep_3drecon.Reconstructor()


def chunk(iterable, chunk_size):
    final_ret = []
    cnt = 0
    ret = []
    for record in iterable:
        if cnt == 0:
            ret = []
        ret.append(record)
        cnt += 1
        if len(ret) == chunk_size:
            final_ret.append(ret)
            ret = []
    if len(final_ret[-1]) != chunk_size:
        final_ret.append(ret)
    return final_ret

# landmark detection in Deep3DRecon
def lm68_2_lm5(in_lm):
    assert in_lm.ndim == 2
    # in_lm: shape=[68,2]
    lm_idx = np.array([31,37,40,43,46,49,55]) - 1
    # 将上述特殊角点的数据取出,得到5个新的角点数据,拼接起来。
    lm = np.stack([in_lm[lm_idx[0],:],np.mean(in_lm[lm_idx[[1,2]],:],0),np.mean(in_lm[lm_idx[[3,4]],:],0),in_lm[lm_idx[5],:],in_lm[lm_idx[6],:]], axis = 0)
    # 将第一个角点放在了第三个位置
    lm = lm[[1,2,0,3,4],:2]
    return lm

def extract_frames_job(fname):
    try:
        out_name=fname.replace(".mp4", "_coeff_pt.npy").replace("/dev/", "/coeff/")
        if os.path.exists(out_name):
            return None
        cap = cv2.VideoCapture(fname)
        frames = []
        while cap.isOpened():
            ret, frame_bgr = cap.read()
            if frame_bgr is None:
                break
            frame_rgb = cv2.cvtColor(frame_bgr, cv2.COLOR_BGR2RGB)
            frames.append(frame_rgb)
        return np.stack(frames)
        # out_name=fname.replace(".mp4", "_coeff_pt.npy").replace("/dev/", "/coeff/")
        # if os.path.exists(out_name):
        #     return None
        # video_reader = VideoReader(fname, ctx=cpu(0))
        # frame_rgb_lst = video_reader.get_batch(list(range(0,len(video_reader)))).asnumpy()
        # return frame_rgb_lst
    except Exception as e:
        print(e)
        return None

def extract_lms_mediapipe_job(frames):
    try:
        if frames is None:
            return None
        with mp_face_mesh.FaceMesh(
                            static_image_mode=False,
                            max_num_faces=1,
                            refine_landmarks=True,
                            min_detection_confidence=0.5) as face_mesh:
            ldms_normed = []
            frame_i = 0
            frame_ids = []
            for i in range(len(frames)):
                # Convert the BGR image to RGB before processing.
                ret = face_mesh.process(frames[i])
                # Print and draw face mesh landmarks on the image.
                if not ret.multi_face_landmarks:
                    print(f"Skip Item: Caught errors when mediapipe get face_mesh, maybe No face detected in some frames!")
                    return None
                else:
                    myFaceLandmarks = []
                    lms = ret.multi_face_landmarks[0]
                    for lm in lms.landmark:
                        myFaceLandmarks.append([lm.x, lm.y, lm.z])
                    ldms_normed.append(myFaceLandmarks)
                frame_ids.append(frame_i)
                frame_i += 1
        bs, H, W, _ = frames.shape
        ldms478 = np.array(ldms_normed)
        lm68 = mediapipe_lm478_to_face_alignment_lm68(ldms478, H, W, return_2d=True)
        lm5_lst = [lm68_2_lm5(lm68[i]) for i in range(lm68.shape[0])]
        lm5 = np.stack(lm5_lst)
        return ldms478, lm68, lm5
    except Exception as e:
        print(e)
        return None
    
def process_video_batch(fname_lst, out_name_lst=None):
    frames_lst = []
    with Timer("load_frames", True):
        for fname in tqdm(fname_lst, desc="decord is loading frames in the batch videos..."):
            res = extract_frames_job(fname)
            frames_lst.append(res)
        # for (i, res) in multiprocess_run_tqdm(extract_frames_job, fname_lst, num_workers=1, desc="decord is loading frames in the batch videos..."):
            # frames_lst.append(res)
    
    lm478s_lst = []
    lm68s_lst = []
    lm5s_lst = []
    with Timer("mediapipe_faceAlign", True):
        # for (i, res) in multiprocess_run_tqdm(extract_lms_mediapipe_job, frames_lst, num_workers=2, desc="mediapipe is predicting face mesh in batch videos..."):
        for i, frames in tqdm(enumerate(frames_lst),total=len(fname_lst), desc="mediapipe is predicting face mesh in batch videos..."):
            res = extract_lms_mediapipe_job(frames)
            if res is None:
                res = (None, None, None)
            lm478s, lm68s, lm5s = res
            lm478s_lst.append(lm478s)
            lm68s_lst.append(lm68s)
            lm5s_lst.append(lm5s)
        
    processed_cnt_in_this_batch = 0
    with Timer("deep_3drecon_pytorch", True):
        for i, fname in tqdm(enumerate(fname_lst), total=len(fname_lst), desc="extracting 3DMM in the batch videos..."):
            video_rgb = frames_lst[i] # [t, 224,224, 3]
            lm478_arr = lm478s_lst[i]
            lm68_arr = lm68s_lst[i]
            lm5_arr = lm5s_lst[i]
            if lm5_arr is None:
                continue
            num_frames = len(video_rgb)
            batch_size = 32
            iter_times = num_frames // batch_size
            last_bs = num_frames % batch_size

            coeff_lst = []
            for i_iter in range(iter_times):
                start_idx = i_iter * batch_size
                batched_images = video_rgb[start_idx: start_idx + batch_size]
                batched_lm5 = lm5_arr[start_idx: start_idx + batch_size]
                coeff, align_img = face_reconstructor.recon_coeff(batched_images, batched_lm5, return_image = True)
                coeff_lst.append(coeff)
            if last_bs != 0:
                batched_images = video_rgb[-last_bs:]
                batched_lm5 = lm5_arr[-last_bs:]
                coeff, align_img = face_reconstructor.recon_coeff(batched_images, batched_lm5, return_image = True)
                coeff_lst.append(coeff)
            coeff_arr = np.concatenate(coeff_lst,axis=0)
            result_dict = {
                'coeff': coeff_arr.reshape([num_frames, -1]).astype(np.float32),
                'lm478': lm478_arr.reshape([num_frames, 478, 3]).astype(np.float32),
                'lm68': lm68_arr.reshape([num_frames, 68, 2]).astype(np.int16),
                'lm5': lm5_arr.reshape([num_frames, 5, 2]).astype(np.int16),
            }
            os.makedirs(os.path.dirname(out_name_lst[i]),exist_ok=True)
            np.save(out_name_lst[i], result_dict)
            processed_cnt_in_this_batch +=1
            
    print(f"In this batch {processed_cnt_in_this_batch} files are processed")



def split_wav(mp4_name):
    try:
        wav_name = mp4_name[:-4] + '.wav'
        if os.path.exists(wav_name):
            return
        video = VideoFileClip(mp4_name,verbose=False)
        dur = video.duration
        audio = video.audio 
        assert audio is not None
        audio.write_audiofile(wav_name,fps=16000,verbose=False,logger=None)
    except Exception as e:
        print(e)
        return None
    
if __name__ == '__main__':
    ### Process Single Long video for NeRF dataset
    # video_id = 'May'
    # video_fname = f"data/raw/videos/{video_id}.mp4"
    # out_fname = f"data/processed/videos/{video_id}/coeff.npy"
    # process_video(video_fname, out_fname)

    ### Process short video clips for LRS3 dataset
    import random 

    from argparse import ArgumentParser
    parser = ArgumentParser()
    parser.add_argument('--lrs3_path', type=str, default='/mnt/sda/yezhenhui/datasets/voxceleb2', help='')
    parser.add_argument('--process_id', type=int, default=0, help='')
    parser.add_argument('--total_process', type=int, default=1, help='')
    args = parser.parse_args()

    import os, glob
    lrs3_dir = args.lrs3_path
    mp4_name_pattern = os.path.join(lrs3_dir, "dev/id*/*/*.mp4")
    mp4_names = glob.glob(mp4_name_pattern)

    if args.total_process > 1:
        assert args.process_id <= args.total_process-1
        num_samples_per_process = len(mp4_names) // args.total_process
        if args.process_id == args.total_process-1:
            mp4_names = mp4_names[args.process_id * num_samples_per_process : ]
        else:
            mp4_names = mp4_names[args.process_id * num_samples_per_process : (args.process_id+1) * num_samples_per_process]
    random.seed(111)
    random.shuffle(mp4_names)
    batched_mp4_names_lst = chunk(mp4_names, chunk_size=1)
    for batch_mp4_names in tqdm(batched_mp4_names_lst, desc='[ROOT]: extracting face mesh and 3DMM in batches...'):
        try:
            for mp4_name in batch_mp4_names:
                split_wav(mp4_name)
            out_names = [mp4_name.replace(".mp4", "_coeff_pt.npy").replace("/dev/", "/coeff/") for mp4_name in batch_mp4_names]
            process_video_batch(batch_mp4_names, out_names)
        except Exception as e:
            print(e)
            continue