|
|
|
|
|
import os
|
|
|
import cv2
|
|
|
import json
|
|
|
import time
|
|
|
import decord
|
|
|
import einops
|
|
|
import librosa
|
|
|
import torch
|
|
|
import random
|
|
|
import argparse
|
|
|
import traceback
|
|
|
import numpy as np
|
|
|
from tqdm import tqdm
|
|
|
from PIL import Image
|
|
|
from einops import rearrange
|
|
|
|
|
|
|
|
|
|
|
|
def get_facemask(ref_image, align_instance, area=1.25):
|
|
|
|
|
|
bsz, f, c, h, w = ref_image.shape
|
|
|
images = rearrange(ref_image, "b f c h w -> (b f) h w c").data.cpu().numpy().astype(np.uint8)
|
|
|
face_masks = []
|
|
|
for image in images:
|
|
|
image_pil = Image.fromarray(image).convert("RGB")
|
|
|
_, _, bboxes_list = align_instance(np.array(image_pil)[:,:,[2,1,0]], maxface=True)
|
|
|
try:
|
|
|
bboxSrc = bboxes_list[0]
|
|
|
except:
|
|
|
bboxSrc = [0, 0, w, h]
|
|
|
x1, y1, ww, hh = bboxSrc
|
|
|
x2, y2 = x1 + ww, y1 + hh
|
|
|
ww, hh = (x2-x1) * area, (y2-y1) * area
|
|
|
center = [(x2+x1)//2, (y2+y1)//2]
|
|
|
x1 = max(center[0] - ww//2, 0)
|
|
|
y1 = max(center[1] - hh//2, 0)
|
|
|
x2 = min(center[0] + ww//2, w)
|
|
|
y2 = min(center[1] + hh//2, h)
|
|
|
|
|
|
face_mask = np.zeros_like(np.array(image_pil))
|
|
|
face_mask[int(y1):int(y2), int(x1):int(x2)] = 1.0
|
|
|
face_masks.append(torch.from_numpy(face_mask[...,:1]))
|
|
|
face_masks = torch.stack(face_masks, dim=0)
|
|
|
face_masks = rearrange(face_masks, "(b f) h w c -> b c f h w", b=bsz, f=f)
|
|
|
face_masks = face_masks.to(device=ref_image.device, dtype=ref_image.dtype)
|
|
|
return face_masks
|
|
|
|
|
|
|
|
|
def encode_audio(wav2vec, audio_feats, fps, num_frames=129):
|
|
|
if fps == 25:
|
|
|
start_ts = [0]
|
|
|
step_ts = [1]
|
|
|
elif fps == 12.5:
|
|
|
start_ts = [0]
|
|
|
step_ts = [2]
|
|
|
else:
|
|
|
start_ts = [0]
|
|
|
step_ts = [1]
|
|
|
|
|
|
num_frames = min(num_frames, 400)
|
|
|
audio_feats = wav2vec.encoder(audio_feats.unsqueeze(0)[:, :, :3000], output_hidden_states=True).hidden_states
|
|
|
audio_feats = torch.stack(audio_feats, dim=2)
|
|
|
audio_feats = torch.cat([torch.zeros_like(audio_feats[:,:4]), audio_feats], 1)
|
|
|
|
|
|
audio_prompts = []
|
|
|
for bb in range(1):
|
|
|
audio_feats_list = []
|
|
|
for f in range(num_frames):
|
|
|
cur_t = (start_ts[bb] + f * step_ts[bb]) * 2
|
|
|
audio_clip = audio_feats[bb:bb+1, cur_t: cur_t+10]
|
|
|
audio_feats_list.append(audio_clip)
|
|
|
audio_feats_list = torch.stack(audio_feats_list, 1)
|
|
|
audio_prompts.append(audio_feats_list)
|
|
|
audio_prompts = torch.cat(audio_prompts)
|
|
|
return audio_prompts |
