前段時間火爆的“螞蟻呀嘿”，將一個人的說話動作和表情遷移到另一張靜態圖中，讓靜態圖中的人臉做出指定動作表情，主要基於FOMM（First Order Motion model）技術。這已經是2年前的技術了，在一些場景中生成的效果並不理想。近期，清華大學團隊在CVPR2022發布最新表情動作遷移論文Thin-Plate Spline Motion Model for Image Animation。本文不具體講論文原理，而是直接將其開源的模型down下來用。效果如下：
第1張圖是靜態照片，第二張是gif驅動動畫，第三張是生成的結果。
本文目的：將開源模型打包封裝成一個單獨接口，讀者只需簡單傳入一張圖片和一個動畫（gif或mp4），即可生成表情遷移動畫（mp4）。
讀者需要安裝好pytorch環境，可前往https://pytorch.org/get-started/locally/根據實際的硬件環境，選擇GPU或cpu版本。
安裝imageio-ffmpeg庫，用於讀取mp4文件。
讀者可以直接跳到最後，獲取源碼，下載源碼資源包後，將文件夾中的對應圖片替換成自己的圖片，即可一鍵生成“螞蟻呀嘿”。
將模型導出為pt後，再創建Model類，將接口封裝到infer函數， 具體代碼如下：
class Model():
def __init__(self, kp="models/kp.pt", aio="models/aio.pt",
device=torch.device('cpu')):
self.device = device
self.kp = torch.jit.load(kp, map_location=device).eval()
self.aio = torch.jit.load(aio, map_location=device).eval()
def relative_kp(self, kp_source, kp_driving, kp_driving_initial):
source_area = ConvexHull(kp_source[0].data.cpu().numpy()).volume
driving_area = ConvexHull(
kp_driving_initial[0].data.cpu().numpy()).volume
adapt_movement_scale = np.sqrt(source_area) / np.sqrt(driving_area)
kp_new = kp_driving
kp_value_diff = (kp_driving - kp_driving_initial)
kp_value_diff *= adapt_movement_scale
kp_new = kp_value_diff + kp_source
return kp_new
def get_kp(self, src):
src = np.expand_dims(src, 0).transpose(0, 3, 1, 2)
src = torch.from_numpy(src).float().to(self.device)
return self.kp(src)
def infer(self, src, driving, src_kp, init_kp):
src = np.expand_dims(src, 0).transpose(0, 3, 1, 2)
src = torch.from_numpy(src).float().to(self.device)
driving = np.expand_dims(driving, 0).transpose(0, 3, 1, 2)
driving = torch.from_numpy(driving).float().to(self.device)
kp_driving = self.kp(driving)
kp_norm = self.relative_kp(kp_source=src_kp,
kp_driving=kp_driving,
kp_driving_initial=init_kp)
with torch.no_grad():
out = self.aio(src, src_kp, kp_norm)
out = out[0].cpu().numpy()
out = out.transpose(1, 2, 0)
return out
其中，get_kp函數用於獲取臉部關鍵點數據。infer函數中，src表示靜態圖，driving表示動態圖中的某一幀，src_kp表示靜態圖的關鍵點，init_kp表示動態圖中的第一幀的關鍵點。
整個調用流程可以拆分為4步：創建模型對象、讀取動圖的每一幀、調用模型、生成幀導出mp4。
前面定義好了Model對象，需要根據GPU和CPU環境，由讀者指定使用具體的pytorch版本，具體代碼如下所示。
def create_model(use_gpu):
if use_gpu:
device = torch.device('cuda')
else:
device = torch.device('cpu')
model = Model(device=device)
return model
上面代碼中，use_gpu是個boolean類型，用於判斷是否使用GPU版本，讀者根據自己的時間情況設置。
調用imageio-ffmpeg庫，讀取mp4或gif文件中的每一幀。具體代碼如下所示，函數返回列表，列表內容為視頻幀：
def read_mp4_or_gif(path):
reader = imageio.get_reader(path)
if path.lower().endswith('.mp4'):
fps = reader.get_meta_data().get('fps')
elif path.lower().endswith('.gif'):
fps = 1000 / Image.open(path).info['duration']
driving_video = []
try:
for im in reader:
im = resize(im, (256, 256))[..., :3]
driving_video.append(im)
except RuntimeError:
pass
reader.close()
return driving_video, fps
因為模型的約束，這裡將每一幀resize到256*256。
模型調用非常簡單，只需讀取靜態圖和動態圖的每一幀，並對靜態圖和動態圖第一幀調用Model類的get_kp函數獲取關鍵點即可。遍歷動態圖的每一幀，將動圖幀、靜態圖、靜態圖關鍵點、動圖第一幀關鍵點一起傳給Model的infer函數即可得到生成的幀。具體代碼如下所示。
def run(use_gpu, src_path, driving_path):
src = imageio.imread(src_path)
src = resize(src, (256, 256))[..., :3]
driving_video, fps = read_mp4_or_gif(driving_path)
model = create_model(use_gpu)
src_kp = model.get_kp(src)
init_kp = model.get_kp(driving_video[0])
outs = []
for driving in driving_video:
out = model.infer(src, driving, src_kp, init_kp)
out = img_as_ubyte(out)
outs.append(out)
return outs, fps
這裡繼續調用imageio-ffmpeg庫，將視頻幀組裝成mp4文件，代碼如下：
def write_mp4(out_path, frames, fps):
imageio.mimsave(out_path, frames, fps=fps)
整個調用流水線如下：
src_path = 'assets/source.png'
driving_path = 'assets/driving2.gif'
frames, fps = run(True, src_path, driving_path)
write_mp4("out.mp4", frames, fps)
Python學習實戰表情遷移，獲取完整源碼。如果您覺得本文有幫助，辛苦您點個不需花錢的贊，您的舉手之勞將對我提供了無限的寫作動力！ 也歡迎關注我的公眾號：Python學習實戰， 第一時間獲取最新文章。
