File size: 2,577 Bytes
872b1a7 | 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 | import sys
import os
# 获取项目根目录并添加到 sys.path 最前面，确保导入正确的 utils 模块
_SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
_PROJECT_ROOT = os.path.abspath(os.path.join(_SCRIPT_DIR, '..', '..'))
if _PROJECT_ROOT not in sys.path:
sys.path.insert(0, _PROJECT_ROOT)
import numpy as np
import torch
from PIL import Image
import torchvision.transforms as T
from omegaconf import OmegaConf
import fire
def init_fn(config_path):
from utils import instantiate
transform = T.Compose([T.Resize((512, 512)), T.ToTensor(), T.Normalize([0.5], [0.5])])
config = OmegaConf.load(config_path)
module = instantiate(config.model, instantiate_module=False)
model = module(config=config)
checkpoint = torch.load(config.resume_ckpt, map_location="cpu")
model.load_state_dict(checkpoint["state_dict"], strict=False)
model.eval()
motion_encoder = model.motion_encoder
return {"transform": transform, "motion_encoder": motion_encoder}
def extract_motion_latent(
mask_image_path='./test_case/test_img_masked.png',
config_path='./configs/head_animator_best_0506.yaml',
save_npz_path='./test_case/test_img_resize.npz',
version="0506"):
sys.path.insert(0, f'./utils/model_{version}')
config_path = config_path.replace("0506", version)
context = init_fn(config_path)
transform = context["transform"]
motion_encoder = context["motion_encoder"]
img = Image.open(mask_image_path).convert("RGB")
img_tensor = transform(img).unsqueeze(0)
with torch.no_grad():
latent = motion_encoder(img_tensor)[0] # [1, 512]
latent_np = latent.numpy()
# 如果文件已存在，先加载原有数据
if os.path.exists(save_npz_path):
existing_data = np.load(save_npz_path, allow_pickle=True)
data_dict = dict(existing_data)
existing_data.close() # 关闭文件
else:
data_dict = {}
# 更新或添加新的键值对
data_dict.update({
'video_id': os.path.basename(save_npz_path)[:-4],
'mask_img_path': mask_image_path,
'ref_img_path': save_npz_path.replace('npz', 'png'),
'motion_latent': latent_np
})
# 保存更新后的数据
np.savez(save_npz_path, **data_dict)
# np.savez(
# save_npz_path,
# video_id=os.path.basename(save_npz_path)[:-4],
# mask_img_path=mask_image_path,
# ref_img_path=save_npz_path.replace('npz', 'png'),
# motion_latent=latent_np
# )
if __name__ == '__main__':
fire.Fire(extract_motion_latent)
|