写的有点乱,后面有时间再整理吧。
问题概览
在做开始,参考了官网和博客的过程,但是过程略微坎坷,问题主要可以总结为:其代码使用了Tensorflow和Pytorch两个框架,我对Tensorflow不是很熟悉,所以直接按照官网的environments.yml安装了,其中的Tensorflow版本使用的Cuda版本为10(我过程中使用的是10.2).而在实际运算的时候我的电脑显卡3060并不支持Cuda10.2.
在官网的数据预处理的命令:
bash process_data.sh Obama
我无法一次性运行该代码,在官网的issue中可以分开执行
python data_util/process_data.py --id=$1 --step=$2
$1这里取得是Obama(视频名) $2取得是步数(0 1 2 ....)
在步数0 1时我使用的是环境为
- Cuda10.2
- Cudnn 7.8.5
- pytorch3D pytorch3d-0.6.1-py37_cu102_pyt181.tar.bz2
- torch==1.8.1+cu102
- torchvision==0.9.1+cu102
在执行到第2步时出现如下的报错 GeForce RTX 3060 Laptop GPU with CUDA capability sm_86 is not compatible with the current PyTorch :
GeForce RTX 3060 Laptop GPU with CUDA capability sm_86 is not compatible with the current PyTorch installation.
The current PyTorch install supports CUDA capabilities sm_37 sm_50 sm_60 sm_61 sm_70 sm_75 compute_37.
If you want to use the GeForce RTX 3060 Laptop GPU GPU with PyTorch, please check the instructions at https://pytorch.org/get-started/locally/
这是个巨坑,30系显卡暂时不支持CUDA11以下版本,CUDA不支持当前显卡的算力。
在此之后,我环境替换为了
- Cuda11.2
- Cudnn8.1.1
- Pytorch3D pytorch3d-0.6.1-py37_cu111_pyt181.tar.bz2
- torch==1.8.1+cu111
- torchvision==0.9.1+cu111
过程参考
博客AD-NeRF 由音频和人脸图像合成人脸视频并表现出自然的说话风格_西西弗Sisyphus的博客-CSDN博客
和 官网GitHub - YudongGuo/AD-NeRF: This repository contains a PyTorch implementation of "AD-NeRF: Audio Driven Neural Radiance Fields for Talking Head Synthesis".
一. PyTorch 环境配置
1.1 PyTorch 环境配置
在这个 environment.yml安装之前需要更换镜像源.这里我参考了下面的博客,解决了我的问题.
Conda 替换镜像源方法尽头,再也不用到处搜镜像源地址_conda换源_五阿哥爱跳舞的博客-CSDN博客在
在安装其中一个torchvision 0.8.1时一直报错 Solving environment: failed with initial frozen solve. Retrying with flexible solve
我参考了这篇博客
使用conda时出现Solving environment: failed with initial frozen solve. Retrying with flexible solve错误_solving environment: failed with initial 会导致下载失败吗?_菜就多练,不要多想的博客-CSDN博客
不过我还是没能用conda命令安装,博客最后的pip install 解决了我的问题. (事实上这也是pytorch官网的安装方法)
pip install torch==1.8.1+cu102 torchvision==0.9.1+cu102 torchaudio===0.8.1 -f https://download.pytorch.org/whl/torch_stable.html
额由于我的显卡为30系的,后面呃又有其他的bug.
pip install torch==1.8.1+cu111 torchvision==0.9.1+cu111 torchaudio==0.8.1 -f https://download.pytorch.org/whl/torch_stable.html
1.2 PyTorch3D 安装
显然我并没有成功地安装,如下
我采用了 下载自己需要的版本的tar.bz2文件 下载地址如下:
唉我的30系显卡无法使用cuda10.2,后面还有一大堆糟糕的事情.
https://anaconda.org/pytorch3d/pytorch3d/files?sort=distribution_type&sort_order=desc&version=&page=7
使用conda 命令安装
conda install pytorch3d-0.6.1-py37_cu102_pyt181.tar.bz2
1.3 下载原始BFM模型
Morphace
填写信息,然后得到邮箱进行下载.
解压后的文件如下:
将"01_MorphableModel.mat" 放到 data_util/face_tracking/3DMM/ 然后
cd data_util/face_tracking 运行python convert_BFM.py
1.4 制作数据集
代码自动会自动找到vids下的文件,当然这一个命令可以分为多部python运行.
具体可以参考AD-NeRF 数据集的制作_ad-nerf 图像边界_西西弗Sisyphus的博客-CSDN博客
然而出现了大面积的错误:具体原因就是我的cuda为11.7而代码里的Tensorflow使用的版本为10
我的解决思路是系统同时有两个cuda,通过环境变量的切换来使用不同的cuda
export CUDA_HOME=/usr/local/cuda-10.2
export LD_LIBRARY_PATH=${CUDA_HOME}/lib64
export PATH=${CUDA_HOME}/bin:${PATH}
source ~/.bashrc
我参考了该博主的解决思路如下:
ubuntu的cuda10和cuda11共存_Dyson Sun的博客-CSDN博客
Ubuntu20.04安装cuda10.2和cudnn7.6.5_ubuntu20安装cuda10.2_小小小白~~的博客-CSDN博客
Ubuntu 20.04 CUDA&cuDNN安装方法(图文教程)-腾讯云开发者社区-腾讯云
最终更改环境变量,更新为cuda10.2
cudnn的安装,网上一搜一大把
最后将cudnn的文件里内容依次复制到对应的include和lib64中.
sudo cp include/cudnn*.h /usr/local/cuda/include
sudo cp -P lib/libcudnn* /usr/local/cuda/lib64
#将include和lib下的文件复制到cuda的安装目录
sudo chmod a+r /usr/local/cuda/include/cudnn*.h /usr/local/cuda/lib64/libcudnn*
cat /usr/local/cuda/include/cudnn_version.h | grep CUDNN_MAJOR -A 2
检验一下,这部因人而异,因为我在安装低版本的时候压根没有cudnn_version.h
使用的命令为
cat /usr/local/cuda-10.2/include/cudnn.h | grep CUDNN_MAJOR -A 2
第0步 提取声音特征
python data_util/process_data.py --id=demosza --step=0
通过ffmpeg命令提取视频中音频数据,存储路径是dataset/$id/aud.wav
通过deepspeech获取音频特征存储路径是dataset/$id/aud.npy
第1步 原始视频转换成图像
python data_util/process_data.py --id=demosza --step=1
将原始视频变成一帧帧的图像,数据存储在dataset/$id/ori_imgs文件夹中
第2步 人脸关键点检测
python data_util/process_data.py --id=demosza --step=2
从dataset/$id/ori_imgs文件夹中读取每一张图像,检测关键点,然后同名存储到dataset/$id/ori_imgs文件夹中,文件扩展名是lms
运行时间几十分钟,运行结果如下:
从dataset/Obama/ori_imgs文件夹中读取每一张图像,检测关键点,然后同名存储到dataset/Obama/ori_imgs文件夹中,文件扩展名是lms:
第3步 人脸解析
python data_util/process_data.py --id=demosza --step=3
输入数据路径是dataset/Obama/ori_imgs
输出数据路径是dataset/Obama/parsing
运行结果如下:
第4步 提取背景图片
python data_util/process_data.py --id=demosza --step=4
这里在第二步人脸关键点检测的时候,确定了最大有效图像个数,默认是100000
并且获取了图像的高度和宽度,后面的操作就操作这些有效的图像。
输出的背景文件是dataset/Obama/bc.jpg
不知道为啥这个运行也很慢,大约一个小时,运行结果如下:
第5步 存储训练数据
python data_util/process_data.py --id=demosza --step=5
通过第4步得到背景数据结合第3步得到解析数据最终得到两种训练数据
运行结果如下:
一个是只有头部和背景的数据,dataset/Obama/head_imgs
另一个是背景和整个上半身的数据,dataset/Obama/com_imgs(CSDN好像不给放obama的图片,会违规)
第6步 估计头部姿态
python data_util/process_data.py --id=demosza --step=6
生成文件的路径是dataset/$id/debug/debug_render
参数存储文件是dataset/$id/track_params.pt
运行结果如下:
这个蛮久的
python data_util/process_data.py --id=Obama --step=7就是保存一下结果
输出包括 HeadNeRF_config.txt TorsoNeRF_config.txt TorsoNeRFTest_config.txt transforms_train.json transforms_val.json
1.5 训练Head-NeRF
训练Head-NeRF
python NeRFs/HeadNeRF/run_nerf.py --config dataset/Obama/HeadNeRF_config.txt
从AD-NeRF/dataset/Obama/logs/Obama_head找到最新的模型, 例如030000_head.tar重命名为head.tar
将head.tar放到AD-NeRF/dataset/Obama/logs/Obama_com中
执行命令
python NeRFs/TorsoNeRF/run_nerf.py --config dataset/Obama/TorsoNeRF_config.txt
1.6 Run AD-NeRF for rendering
两个训练完了就可以用自己的音频特征去驱动之前的目标
python NeRFs/TorsoNeRF/run_nerf.py --config dataset/$id/TorsoNeRFTest_config.txt --aud_file=${deepspeechfile.npy} --test_size=-1
这里输入自己的音频特征文件,可以使用deepspeech获取音频特征文件npy,制作数据集的第一部就是这个.不过我们只需要npy文件,
而步骤1是提取视频中的音频文件,再提取音频特征.
我们只需要执行这里面的extract_ds_features.py文件即可,输入我们自己的wav文件.
得到自己的音频特征文件npy,就可以进行语音驱动face了
python NeRFs/TorsoNeRF/run_nerf.py --config dataset/$id/TorsoNeRFTest_config.txt --aud_file=${deepspeechfile.npy} --test_size=-1
结果在:/home/sza/AD-NeRF/dataset/Obama/logs/Obama_com/test_aud_rst中的result.avi(没有声音)
最后通过ffmpeg合并音频和视频.
ffmpeg -i video.mp4/video.avi -i audio.wav -c:v copy -c:a aac -map 0:v:0 -map 1:a:0 output.mp4
音频: /home/sza/AD-NeRF/sza/output.wav
视频: /home/sza/AD-NeRF/dataset/Obama/logs/Obama_com/test_aud_rst/result.avi
在这个命令中,video.mp4
是你的视频文件,audio.wav
是你的音频文件,output.mp4
是输出的合并文件名。
命令解释:
-
-i video.mp4
:指定输入视频文件。 -
-i audio.wav
:指定输入音频文件。 -
-c:v copy
:表示对视频流进行复制,保持原始编码格式,避免重新编码。 -
-c:a aac
:表示对音频流进行AAC编码,用于生成合并后的视频文件。 -
-map 0:v:0
:选择输入文件的第一个视频流。 -
-map 1:a:0
:选择输入文件的第一个音频流。 -
output.mp4
:指定输出文件名。
运行命令后,FFmpeg将合并音频和视频,并生成一个新的视频文件output.mp4
。请确保你已经安装了FFmpeg,并将其添加到系统的环境变量中,这样你就可以在命令行中直接使用FFmpeg命令。最后的结果:
二 py文件概览
2.1 process_data.py
import cv2
import numpy as np
import face_alignment
from skimage import io
import torch
import torch.nn.functional as F
import json
import os
from sklearn.neighbors import NearestNeighbors
from pathlib import Path
import argparse
def euler2rot(euler_angle):
batch_size = euler_angle.shape[0]
theta = euler_angle[:, 0].reshape(-1, 1, 1)
phi = euler_angle[:, 1].reshape(-1, 1, 1)
psi = euler_angle[:, 2].reshape(-1, 1, 1)
one = torch.ones((batch_size, 1, 1), dtype=torch.float32,
device=euler_angle.device)
zero = torch.zeros((batch_size, 1, 1), dtype=torch.float32,
device=euler_angle.device)
rot_x = torch.cat((
torch.cat((one, zero, zero), 1),
torch.cat((zero, theta.cos(), theta.sin()), 1),
torch.cat((zero, -theta.sin(), theta.cos()), 1),
), 2)
rot_y = torch.cat((
torch.cat((phi.cos(), zero, -phi.sin()), 1),
torch.cat((zero, one, zero), 1),
torch.cat((phi.sin(), zero, phi.cos()), 1),
), 2)
rot_z = torch.cat((
torch.cat((psi.cos(), -psi.sin(), zero), 1),
torch.cat((psi.sin(), psi.cos(), zero), 1),
torch.cat((zero, zero, one), 1)
), 2)
return torch.bmm(rot_x, torch.bmm(rot_y, rot_z))
parser = argparse.ArgumentParser()
parser.add_argument('--id', type=str,
default='obama', help='identity of target person')
parser.add_argument('--step', type=int,
default=0, help='step for running')
args = parser.parse_args()
id = args.id
vid_file = os.path.join('dataset', 'vids', id+'.mp4')
if not os.path.isfile(vid_file):
print('no video')
exit()
id_dir = os.path.join('dataset', id)
Path(id_dir).mkdir(parents=True, exist_ok=True)
ori_imgs_dir = os.path.join('dataset', id, 'ori_imgs')
Path(ori_imgs_dir).mkdir(parents=True, exist_ok=True)
parsing_dir = os.path.join(id_dir, 'parsing')
Path(parsing_dir).mkdir(parents=True, exist_ok=True)
head_imgs_dir = os.path.join('dataset', id, 'head_imgs')
Path(head_imgs_dir).mkdir(parents=True, exist_ok=True)
com_imgs_dir = os.path.join('dataset', id, 'com_imgs')
Path(com_imgs_dir).mkdir(parents=True, exist_ok=True)
running_step = args.step
# # Step 0: extract wav & deepspeech feature, better run in terminal to parallel with
# below commands since this may take a few minutes
'''
步骤 0:提取 wav 和 deepspeech 特征,最好在终端中运行以与以下命令并行,因为这可能需要几分钟
通过ffmpeg命令提取视频中音频数据,存储路径是dataset/Obama/aud.wav
通过deepspeech获取音频特征存储路径是dataset/Obama/aud.npy
原文中提到的分辨率是450×450,不过看代码应该没关系,帧速率是25fps,音频采样率是16kHz
'''
if running_step == 0:
print('--- Step0: extract deepspeech feature ---')
wav_file = os.path.join(id_dir, 'aud.wav')
extract_wav_cmd = 'ffmpeg -i ' + vid_file + ' -f wav -ar 16000 ' + wav_file
os.system(extract_wav_cmd)
extract_ds_cmd = 'python data_util/deepspeech_features/extract_ds_features.py --input=' + id_dir
os.system(extract_ds_cmd)
exit()
# Step 1: extract images
'''
# 第1步:提取图像
将原始视频变成一帧帧的图像,数据存储在dataset/Obama/ori_imgs文件夹中
'''
if running_step == 1:
print('--- Step1: extract images from vids ---')
cap = cv2.VideoCapture(vid_file)
frame_num = 0
while(True):
_, frame = cap.read()
if frame is None:
break
cv2.imwrite(os.path.join(ori_imgs_dir, str(frame_num) + '.jpg'), frame)
frame_num = frame_num + 1
cap.release()
exit()
# Step 2: detect lands
'''
第步 人脸关键点检测
https://github.com/1adrianb/face-alignment
从dataset/Obama/ori_imgs文件夹中读取每一张图像,检测关键点,然后同名存储到dataset/Obama/ori_imgs文件夹中,文件扩展名是lms
用的是face-alignment,face-alignment可以检测2D或者3D的人脸关键点
'''
if running_step == 2:
print('--- Step 2: detect landmarks ---')
fa = face_alignment.FaceAlignment(
face_alignment.LandmarksType._2D, flip_input=False)
for image_path in os.listdir(ori_imgs_dir):
if image_path.endswith('.jpg'):
input = io.imread(os.path.join(ori_imgs_dir, image_path))[:, :, :3]
preds = fa.get_landmarks(input)
if len(preds) > 0:
lands = preds[0].reshape(-1, 2)[:,:2]
np.savetxt(os.path.join(ori_imgs_dir, image_path[:-3] + 'lms'), lands, '%f')
max_frame_num = 100000
valid_img_ids = []
for i in range(max_frame_num):
if os.path.isfile(os.path.join(ori_imgs_dir, str(i) + '.lms')):
valid_img_ids.append(i)
valid_img_num = len(valid_img_ids)
tmp_img = cv2.imread(os.path.join(ori_imgs_dir, str(valid_img_ids[0])+'.jpg'))
h, w = tmp_img.shape[0], tmp_img.shape[1]
# Step 3: face parsing
'''
第3步 人脸解析
输入数据路径是dataset/Obama/ori_imgs
输出数据路径是dataset/Obama/parsing
'''
if running_step == 3:
print('--- Step 3: face parsing ---')
face_parsing_cmd = 'python data_util/face_parsing/test.py --respath=dataset/' + \
id + '/parsing --imgpath=dataset/' + id + '/ori_imgs'
os.system(face_parsing_cmd)
# Step 4: extract bc image
'''
第4步 提取背景图片
这里在第二步人脸关键点检测的时候,确定了最大有效图像个数,默认是100000
并且获取了图像的高度和宽度,后面的操作就操作这些有效的图像。
输出的背景文件是dataset/Obama/bc.jpg
'''
if running_step == 4:
print('--- Step 4: extract background image ---')
sel_ids = np.array(valid_img_ids)[np.arange(0, valid_img_num, 20)]
all_xys = np.mgrid[0:h, 0:w].reshape(2, -1).transpose()
distss = []
for i in sel_ids:
parse_img = cv2.imread(os.path.join(id_dir, 'parsing', str(i) + '.png'))
bg = (parse_img[..., 0] == 255) & (
parse_img[..., 1] == 255) & (parse_img[..., 2] == 255)
fg_xys = np.stack(np.nonzero(~bg)).transpose(1, 0)
nbrs = NearestNeighbors(n_neighbors=1, algorithm='kd_tree').fit(fg_xys)
dists, _ = nbrs.kneighbors(all_xys)
distss.append(dists)
distss = np.stack(distss)
print(distss.shape)
max_dist = np.max(distss, 0)
max_id = np.argmax(distss, 0)
bc_pixs = max_dist > 5
bc_pixs_id = np.nonzero(bc_pixs)
bc_ids = max_id[bc_pixs]
imgs = []
num_pixs = distss.shape[1]
for i in sel_ids:
img = cv2.imread(os.path.join(ori_imgs_dir, str(i) + '.jpg'))
imgs.append(img)
imgs = np.stack(imgs).reshape(-1, num_pixs, 3)
bc_img = np.zeros((h*w, 3), dtype=np.uint8)
bc_img[bc_pixs_id, :] = imgs[bc_ids, bc_pixs_id, :]
bc_img = bc_img.reshape(h, w, 3)
max_dist = max_dist.reshape(h, w)
bc_pixs = max_dist > 5
bg_xys = np.stack(np.nonzero(~bc_pixs)).transpose()
fg_xys = np.stack(np.nonzero(bc_pixs)).transpose()
nbrs = NearestNeighbors(n_neighbors=1, algorithm='kd_tree').fit(fg_xys)
distances, indices = nbrs.kneighbors(bg_xys)
bg_fg_xys = fg_xys[indices[:, 0]]
print(fg_xys.shape)
print(np.max(bg_fg_xys), np.min(bg_fg_xys))
bc_img[bg_xys[:, 0], bg_xys[:, 1],
:] = bc_img[bg_fg_xys[:, 0], bg_fg_xys[:, 1], :]
cv2.imwrite(os.path.join(id_dir, 'bc.jpg'), bc_img)
# Step 5: save training images
'''
第5步 存储训练数据
通过第4步得到背景数据结合第3步得到解析数据最终得到两种训练数据
一个是只有头部和背景的数据,dataset/Obama/head_imgs
另一个是背景和整个上半身的数据,dataset/Obama/com_imgs
'''
if running_step == 5:
print('--- Step 5: save training images ---')
bc_img = cv2.imread(os.path.join(id_dir, 'bc.jpg'))
for i in valid_img_ids:
parsing_img = cv2.imread(os.path.join(parsing_dir, str(i) + '.png'))
head_part = (parsing_img[:, :, 0] == 255) & (
parsing_img[:, :, 1] == 0) & (parsing_img[:, :, 2] == 0)
bc_part = (parsing_img[:, :, 0] == 255) & (
parsing_img[:, :, 1] == 255) & (parsing_img[:, :, 2] == 255)
img = cv2.imread(os.path.join(ori_imgs_dir, str(i) + '.jpg'))
img[bc_part] = bc_img[bc_part]
cv2.imwrite(os.path.join(com_imgs_dir, str(i) + '.jpg'), img)
img[~head_part] = bc_img[~head_part]
cv2.imwrite(os.path.join(head_imgs_dir, str(i) + '.jpg'), img)
# Step 6: estimate head pose
'''
第6步 估计头部姿态
生成文件的路径是dataset/Obama/debug/debug_render
参数存储文件是dataset/Obama/track_params.pt
'''
if running_step == 6:
print('--- Estimate Head Pose ---')
est_pose_cmd = 'python data_util/face_tracking/face_tracker.py --idname=' + \
id + ' --img_h=' + str(h) + ' --img_w=' + str(w) + \
' --frame_num=' + str(max_frame_num)
os.system(est_pose_cmd)
exit()
# Step 7: save transform param & write config file
'''
第7步 写入Transform参数和配置文件
通过第6步得到track_params.pt参数,输出包括
HeadNeRF_config.txt
TorsoNeRF_config.txt
TorsoNeRFTest_config.txt
transforms_train.json
transforms_val.json
'''
if running_step == 7:
print('--- Step 7: Save Transform Param ---')
'''
加载参数:从指定路径下加载包含转换参数的文件 track_params.pt,并将参数存储在 params_dict 字典中。
'''
params_dict = torch.load(os.path.join(id_dir, 'track_params.pt'))
'''
从params_dict中提取焦距(focal)、欧拉角(euler)和平移(trans)。此外,平移参数会除以10.0进行缩放。
'''
focal_len = params_dict['focal']
euler_angle = params_dict['euler']
trans = params_dict['trans'] / 10.0
'''
计算有效样本数量:获取欧拉角的形状,并计算有效样本的数量。
'''
valid_num = euler_angle.shape[0]
'''
获取欧拉角的有效数量(valid_num),然后计算用于训练和验证集划分的索引(train_val_split)。
其中,训练集的索引范围是从0到train_val_split,验证集的索引范围是从train_val_split到valid_num。
'''
train_val_split = int(valid_num*10/11)
train_ids = torch.arange(0, train_val_split)
val_ids = torch.arange(train_val_split, valid_num)
'''
将欧拉角转换为旋转矩阵(rot),并计算其逆矩阵(rot_inv)。
计算平移的逆变换(trans_inv),通过使用旋转矩阵的逆和平移的负向量的乘积。
'''
rot = euler2rot(euler_angle)
rot_inv = rot.permute(0, 2, 1)
trans_inv = -torch.bmm(rot_inv, trans.unsqueeze(2))
'''
创建一个4x4的单位矩阵(pose),用于存储姿态变换矩阵。
创建两个列表save_ids和train_val_ids,分别用于保存训练集和验证集的ID。
'''
pose = torch.eye(4, dtype=torch.float32)
save_ids = ['train', 'val']
train_val_ids = [train_ids, val_ids]
'''
计算平移参数trans在Z轴上的平均值,并取其相反数作为mean_z。
'''
mean_z = -float(torch.mean(trans[:, 2]).item())
'''
进行一个循环,循环两次,分别处理训练集和验证集。
在循环中,创建一个字典transform_dict,用于存储变换的元数据信息。
将焦距(focal_len)、图像宽度的一半(w/2.0)和图像高度的一半(h/2.0)保存到transform_dict中。
创建一个空列表frames,用于存储每个图像帧的元数据。
获取当前处理的集合的ID范围(ids),以及保存该集合的ID字符串(save_id)。
在ids中进行循环迭代,处理每个图像帧。
获取当前帧的索引i,并创建一个字典frame_dict,用于存储该帧的元数据。
将图像ID(valid_img_ids[i])和音频ID(valid_img_ids[i])保存到frame_dict中。
根据当前帧的索引i,将旋转矩阵和平移参数更新到pose矩阵中,并将变换矩阵(pose)转换为NumPy数组,并将其保存到frame_dict中。
加载当前帧对应的人脸关键点坐标(lms),并根据关键点坐标计算人脸区域的矩形坐标。
将人脸矩形区域(rect)保存到frame_dict中。
将当前帧的元数据(frame_dict)添加到transform_dict['frames']列表中。
将transform_dict以JSON格式写入到文件中,文件名格式为transforms_{save_id}.json。
'''
for i in range(2):
transform_dict = dict()
transform_dict['focal_len'] = float(focal_len[0])
transform_dict['cx'] = float(w/2.0)
transform_dict['cy'] = float(h/2.0)
transform_dict['frames'] = []
ids = train_val_ids[i]
save_id = save_ids[i]
for i in ids:
i = i.item()
frame_dict = dict()
frame_dict['img_id'] = int(valid_img_ids[i])
frame_dict['aud_id'] = int(valid_img_ids[i])
pose[:3, :3] = rot_inv[i]
pose[:3, 3] = trans_inv[i, :, 0]
frame_dict['transform_matrix'] = pose.numpy().tolist()
lms = np.loadtxt(os.path.join(
ori_imgs_dir, str(valid_img_ids[i]) + '.lms'))
min_x, max_x = np.min(lms, 0)[0], np.max(lms, 0)[0]
cx = int((min_x+max_x)/2.0)
cy = int(lms[27, 1])
h_w = int((max_x-cx)*1.5)
h_h = int((lms[8, 1]-cy)*1.15)
rect_x = cx - h_w
rect_y = cy - h_h
if rect_x < 0:
rect_x = 0
if rect_y < 0:
rect_y = 0
rect_w = min(w-1-rect_x, 2*h_w)
rect_h = min(h-1-rect_y, 2*h_h)
rect = np.array((rect_x, rect_y, rect_w, rect_h), dtype=np.int32)
frame_dict['face_rect'] = rect.tolist()
transform_dict['frames'].append(frame_dict)
with open(os.path.join(id_dir, 'transforms_' + save_id + '.json'), 'w') as fp:
json.dump(transform_dict, fp, indent=2, separators=(',', ': '))
'''
获取根目录的路径,并构建HeadNeRF的配置文件路径(HeadNeRF_config_file)。
'''
dir_path = os.path.dirname(os.path.dirname(os.path.realpath(__file__)))
testskip = int(val_ids.shape[0]/7)
'''
使用文件操作,将HeadNeRF的配置信息写入到配置文件中。
创建一个名为id + '_head'的文件夹,用于保存HeadNeRF的日志和输出。
'''
HeadNeRF_config_file = os.path.join(id_dir, 'HeadNeRF_config.txt')
with open(HeadNeRF_config_file, 'w') as file:
file.write('expname = ' + id + '_head\n')
file.write('datadir = ' + os.path.join(dir_path, 'dataset', id) + '\n')
file.write('basedir = ' + os.path.join(dir_path,
'dataset', id, 'logs') + '\n')
file.write('near = ' + str(mean_z-0.2) + '\n')
file.write('far = ' + str(mean_z+0.4) + '\n')
file.write('testskip = ' + str(testskip) + '\n')
Path(os.path.join(dir_path, 'dataset', id, 'logs', id + '_head')
).mkdir(parents=True, exist_ok=True)
'''
torso是躯干的意思
构建ComNeRF的配置文件路径(ComNeRF_config_file),并将ComNeRF的配置信息写入到配置文件中。
创建一个名为id + '_com'的文件夹,用于保存ComNeRF的日志和输出。
构建ComNeRFTest的配置文件路径(ComNeRFTest_config_file),并将ComNeRFTest的配置信息写入到配置文件中。
'''
ComNeRF_config_file = os.path.join(id_dir, 'TorsoNeRF_config.txt')
with open(ComNeRF_config_file, 'w') as file:
file.write('expname = ' + id + '_com\n')
file.write('datadir = ' + os.path.join(dir_path, 'dataset', id) + '\n')
file.write('basedir = ' + os.path.join(dir_path,
'dataset', id, 'logs') + '\n')
file.write('near = ' + str(mean_z-0.2) + '\n')
file.write('far = ' + str(mean_z+0.4) + '\n')
file.write('testskip = ' + str(testskip) + '\n')
Path(os.path.join(dir_path, 'dataset', id, 'logs', id + '_com')
).mkdir(parents=True, exist_ok=True)
ComNeRFTest_config_file = os.path.join(id_dir, 'TorsoNeRFTest_config.txt')
with open(ComNeRFTest_config_file, 'w') as file:
file.write('expname = ' + id + '_com\n')
file.write('datadir = ' + os.path.join(dir_path, 'dataset', id) + '\n')
file.write('basedir = ' + os.path.join(dir_path,
'dataset', id, 'logs') + '\n')
file.write('near = ' + str(mean_z-0.2) + '\n')
file.write('far = ' + str(mean_z+0.4) + '\n')
file.write('with_test = ' + str(1) + '\n')
file.write('test_pose_file = transforms_val.json' + '\n')
print(id + ' data processed done!')
2.2 run_nerf.py
关于nerf,这里很多的内容可以看这两篇博客
【代码详解】nerf-pytorch代码逐行分析_nerf_pytorch代码解读csdn_YuhsiHu的博客-CSDN博客文章来源:https://www.toymoban.com/news/detail-772166.html
NeRF神经辐射场学习笔记(二)——Pytorch版NeRF实现以及代码注释_nerf pytorch_右边的口袋的博客-CSDN博客文章来源地址https://www.toymoban.com/news/detail-772166.html
from load_audface import load_audface_data
import os
import sys
import numpy as np
import imageio
import json
import random
import time
import torch
import torch.nn as nn
import torch.nn.functional as F
from tqdm import tqdm, trange
from natsort import natsorted
from run_nerf_helpers import *
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
np.random.seed(0)
DEBUG = False
def batchify(fn, chunk):
"""Constructs a version of 'fn' that applies to smaller batches.
"""
if chunk is None:
return fn
def ret(inputs):
return torch.cat([fn(inputs[i:i+chunk]) for i in range(0, inputs.shape[0], chunk)], 0)
return ret
def run_network(inputs, viewdirs, aud_para, fn, embed_fn, embeddirs_fn, netchunk=1024*64):
"""Prepares inputs and applies network 'fn'.
"""
inputs_flat = torch.reshape(inputs, [-1, inputs.shape[-1]])
embedded = embed_fn(inputs_flat)
aud = aud_para.unsqueeze(0).expand(inputs_flat.shape[0], -1)
embedded = torch.cat((embedded, aud), -1)
if viewdirs is not None:
input_dirs = viewdirs[:, None].expand(inputs.shape)
input_dirs_flat = torch.reshape(input_dirs, [-1, input_dirs.shape[-1]])
embedded_dirs = embeddirs_fn(input_dirs_flat)
embedded = torch.cat([embedded, embedded_dirs], -1)
outputs_flat = batchify(fn, netchunk)(embedded)
outputs = torch.reshape(outputs_flat, list(
inputs.shape[:-1]) + [outputs_flat.shape[-1]])
return outputs
def batchify_rays(rays_flat, bc_rgb, aud_para, chunk=1024*32, **kwargs):
"""Render rays in smaller minibatches to avoid OOM.
"""
all_ret = {}
for i in range(0, rays_flat.shape[0], chunk):
ret = render_rays(rays_flat[i:i+chunk], bc_rgb[i:i+chunk],
aud_para, **kwargs)
for k in ret:
if k not in all_ret:
all_ret[k] = []
all_ret[k].append(ret[k])
all_ret = {k: torch.cat(all_ret[k], 0) for k in all_ret}
return all_ret
def render_dynamic_face(H, W, focal, cx, cy, chunk=1024*32, rays=None, bc_rgb=None, aud_para=None,
c2w=None, ndc=True, near=0., far=1.,
use_viewdirs=False, c2w_staticcam=None,
**kwargs):
if c2w is not None:
# special case to render full image
rays_o, rays_d = get_rays(H, W, focal, c2w, cx, cy)
bc_rgb = bc_rgb.reshape(-1, 3)
else:
# use provided ray batch
rays_o, rays_d = rays
if use_viewdirs:
# provide ray directions as input
viewdirs = rays_d
if c2w_staticcam is not None:
# special case to visualize effect of viewdirs
rays_o, rays_d = get_rays(H, W, focal, c2w_staticcam, cx, cy)
viewdirs = viewdirs / torch.norm(viewdirs, dim=-1, keepdim=True)
viewdirs = torch.reshape(viewdirs, [-1, 3]).float()
sh = rays_d.shape # [..., 3]
if ndc:
# for forward facing scenes
rays_o, rays_d = ndc_rays(H, W, focal, 1., rays_o, rays_d)
# Create ray batch
rays_o = torch.reshape(rays_o, [-1, 3]).float()
rays_d = torch.reshape(rays_d, [-1, 3]).float()
near, far = near * \
torch.ones_like(rays_d[..., :1]), far * \
torch.ones_like(rays_d[..., :1])
rays = torch.cat([rays_o, rays_d, near, far], -1)
if use_viewdirs:
rays = torch.cat([rays, viewdirs], -1)
# Render and reshape
all_ret = batchify_rays(rays, bc_rgb, aud_para, chunk, **kwargs)
for k in all_ret:
k_sh = list(sh[:-1]) + list(all_ret[k].shape[1:])
all_ret[k] = torch.reshape(all_ret[k], k_sh)
k_extract = ['rgb_map', 'disp_map', 'acc_map', 'last_weight']
ret_list = [all_ret[k] for k in k_extract]
ret_dict = {k: all_ret[k] for k in all_ret if k not in k_extract}
return ret_list + [ret_dict]
def render(H, W, focal, cx, cy, chunk=1024*32, rays=None, c2w=None, ndc=True,
near=0., far=1.,
use_viewdirs=False, c2w_staticcam=None,
**kwargs):
"""Render rays
Args:
H: int. Height of image in pixels.
W: int. Width of image in pixels.
focal: float. Focal length of pinhole camera.
chunk: int. Maximum number of rays to process simultaneously. Used to
control maximum memory usage. Does not affect final results.
rays: array of shape [2, batch_size, 3]. Ray origin and direction for
each example in batch.
c2w: array of shape [3, 4]. Camera-to-world transformation matrix.
ndc: bool. If True, represent ray origin, direction in NDC coordinates.
near: float or array of shape [batch_size]. Nearest distance for a ray.
far: float or array of shape [batch_size]. Farthest distance for a ray.
use_viewdirs: bool. If True, use viewing direction of a point in space in model.
c2w_staticcam: array of shape [3, 4]. If not None, use this transformation matrix for
camera while using other c2w argument for viewing directions.
Returns:
rgb_map: [batch_size, 3]. Predicted RGB values for rays.
disp_map: [batch_size]. Disparity map. Inverse of depth.
acc_map: [batch_size]. Accumulated opacity (alpha) along a ray.
extras: dict with everything returned by render_rays().
"""
if c2w is not None:
# special case to render full image
rays_o, rays_d = get_rays(H, W, focal, c2w, cx, cy)
else:
# use provided ray batch
rays_o, rays_d = rays
if use_viewdirs:
# provide ray directions as input
viewdirs = rays_d
if c2w_staticcam is not None:
# special case to visualize effect of viewdirs
rays_o, rays_d = get_rays(H, W, focal, c2w_staticcam, cx, cy)
viewdirs = viewdirs / torch.norm(viewdirs, dim=-1, keepdim=True)
viewdirs = torch.reshape(viewdirs, [-1, 3]).float()
sh = rays_d.shape # [..., 3]
if ndc:
# for forward facing scenes
rays_o, rays_d = ndc_rays(H, W, focal, 1., rays_o, rays_d)
# Create ray batch
rays_o = torch.reshape(rays_o, [-1, 3]).float()
rays_d = torch.reshape(rays_d, [-1, 3]).float()
near, far = near * \
torch.ones_like(rays_d[..., :1]), far * \
torch.ones_like(rays_d[..., :1])
rays = torch.cat([rays_o, rays_d, near, far], -1)
if use_viewdirs:
rays = torch.cat([rays, viewdirs], -1)
# Render and reshape
all_ret = batchify_rays(rays, chunk, **kwargs)
for k in all_ret:
k_sh = list(sh[:-1]) + list(all_ret[k].shape[1:])
all_ret[k] = torch.reshape(all_ret[k], k_sh)
k_extract = ['rgb_map', 'disp_map', 'acc_map']
ret_list = [all_ret[k] for k in k_extract]
ret_dict = {k: all_ret[k] for k in all_ret if k not in k_extract}
return ret_list + [ret_dict]
def render_path(render_poses, aud_paras, bc_img, hwfcxy,
chunk, render_kwargs, gt_imgs=None, savedir=None, render_factor=0):
H, W, focal, cx, cy = hwfcxy
if render_factor != 0:
# Render downsampled for speed
H = H//render_factor
W = W//render_factor
focal = focal/render_factor
rgbs = []
disps = []
last_weights = []
t = time.time()
for i, c2w in enumerate(tqdm(render_poses)):
print(i, time.time() - t)
t = time.time()
rgb, disp, acc, last_weight, _ = render_dynamic_face(
H, W, focal, cx, cy, chunk=chunk, c2w=c2w[:3,
:4], aud_para=aud_paras[i], bc_rgb=bc_img,
**render_kwargs)
rgbs.append(rgb.cpu().numpy())
disps.append(disp.cpu().numpy())
last_weights.append(last_weight.cpu().numpy())
if i == 0:
print(rgb.shape, disp.shape)
"""
if gt_imgs is not None and render_factor==0:
p = -10. * np.log10(np.mean(np.square(rgb.cpu().numpy() - gt_imgs[i])))
print(p)
"""
if savedir is not None:
rgb8 = to8b(rgbs[-1])
filename = os.path.join(savedir, '{:03d}.png'.format(i))
imageio.imwrite(filename, rgb8)
rgbs = np.stack(rgbs, 0)
disps = np.stack(disps, 0)
last_weights = np.stack(last_weights, 0)
return rgbs, disps, last_weights
def create_nerf(args):
"""Instantiate NeRF's MLP model.
函数接受一个参数args,其中包含了NeRF模型的各种配置选项。
"""
# multires = 10 位置编码操作对于3D位置信息的所升维数,默认L=10
# i_embed = 0 是否加入位置编码操作,设为0是默认采用位置编码方法,-1则无
'''
首先,函数调用get_embedder函数来获取位置编码器(embedder)对象和输入通道数input_ch。
位置编码器用于将输入的位置信息进行编码,以供后续的神经网络模型使用。
'''
embed_fn, input_ch = get_embedder(args.multires, args.i_embed)
input_ch_views = 0
embeddirs_fn = None
# use_viewdirs = True用完整的5D信息代替3D信息
if args.use_viewdirs:
'''
接着,根据是否使用视角方向(view directions)的配置选项,
同样使用get_embedder函数获取视角编码器(embedder)对象和输入通道数input_ch_views。
'''
embeddirs_fn, input_ch_views = get_embedder(
args.multires_views, args.i_embed)
output_ch = 5 if args.N_importance > 0 else 4
skips = [4]
'''
根据NeRF模型的其他配置选项,创建一个FaceNeRF模型对象model。
FaceNeRF是NeRF模型的实现,它接受各种输入参数来定义模型的深度、宽度、输入通道数、输出通道数等。
'''
model = FaceNeRF(D=args.netdepth, W=args.netwidth,
input_ch=input_ch, dim_aud=args.dim_aud,
output_ch=output_ch, skips=skips,
input_ch_views=input_ch_views, use_viewdirs=args.use_viewdirs).to(device)
grad_vars = list(model.parameters())
model_fine = None
'''
配置选项中指定了N_importance大于0,则创建一个额外的精细模型model_fine,用于处理更重要的样本
'''
if args.N_importance > 0:
model_fine = FaceNeRF(D=args.netdepth_fine, W=args.netwidth_fine,
input_ch=input_ch, dim_aud=args.dim_aud,
output_ch=output_ch, skips=skips,
input_ch_views=input_ch_views, use_viewdirs=args.use_viewdirs).to(device)
grad_vars += list(model_fine.parameters())
'''
创建了一个network_query_fn函数,用于执行神经网络的前向传播。
这个函数会将输入数据、视角方向、音频参数等传递给网络模型,并返回输出结果。
'''
def network_query_fn(inputs, viewdirs, aud_para, network_fn): \
return run_network(inputs, viewdirs, aud_para, network_fn,
embed_fn=embed_fn, embeddirs_fn=embeddirs_fn, netchunk=args.netchunk)
'''
创建了一个Adam优化器,用于优化模型的参数。
'''
# Create optimizer
optimizer = torch.optim.Adam(
params=grad_vars, lr=args.lrate, betas=(0.9, 0.999))
start = 0
basedir = args.basedir
expname = args.expname
##########################
# Load checkpoints
'''
在加载检查点(checkpoint)时,函数会检查是否指定了检查点文件路径。
如果有,则加载指定的检查点文件,恢复模型和优化器的状态。
'''
if args.ft_path is not None and args.ft_path != 'None':
ckpts = [args.ft_path]
else:
ckpts = [os.path.join(basedir, expname, f) for f in natsorted(
os.listdir(os.path.join(basedir, expname))) if 'tar' in f]
print('Found ckpts', ckpts)
learned_codes_dict = None
AudNet_state = None
AudAttNet_state = None
optimizer_aud_state = None
optimizer_audatt_state = None
if len(ckpts) > 0 and not args.no_reload:
ckpt_path = ckpts[-1]
print('Reloading from', ckpt_path)
ckpt = torch.load(ckpt_path)
start = ckpt['global_step']
optimizer.load_state_dict(ckpt['optimizer_state_dict'])
AudNet_state = ckpt['network_audnet_state_dict']
optimizer_aud_state = ckpt['optimizer_aud_state_dict']
# Load model
model.load_state_dict(ckpt['network_fn_state_dict'])
if model_fine is not None:
model_fine.load_state_dict(ckpt['network_fine_state_dict'])
if 'network_audattnet_state_dict' in ckpt:
AudAttNet_state = ckpt['network_audattnet_state_dict']
if 'optimize_audatt_state_dict' in ckpt:
optimizer_audatt_state = ckpt['optimize_audatt_state_dict']
##########################
'''
最后,函数定义了render_kwargs_train和render_kwargs_test字典,其中包含了渲染模型所需的各种参数配置。
'''
render_kwargs_train = {
'network_query_fn': network_query_fn,
'perturb': args.perturb,
'N_importance': args.N_importance,
'network_fine': model_fine,
'N_samples': args.N_samples,
'network_fn': model,
'use_viewdirs': args.use_viewdirs,
'white_bkgd': args.white_bkgd,
'raw_noise_std': args.raw_noise_std,
}
# NDC only good for LLFF-style forward facing data
if args.dataset_type != 'llff' or args.no_ndc:
print('Not ndc!')
render_kwargs_train['ndc'] = False
render_kwargs_train['lindisp'] = args.lindisp
render_kwargs_test = {
k: render_kwargs_train[k] for k in render_kwargs_train}
render_kwargs_test['perturb'] = False
render_kwargs_test['raw_noise_std'] = 0.
return render_kwargs_train, render_kwargs_test, start, grad_vars, optimizer, learned_codes_dict, \
AudNet_state, optimizer_aud_state, AudAttNet_state, optimizer_audatt_state
def raw2outputs(raw, z_vals, rays_d, bc_rgb, raw_noise_std=0, white_bkgd=False, pytest=False):
"""Transforms model's predictions to semantically meaningful values.
Args:
raw: [num_rays, num_samples along ray, 4]. Prediction from model.
z_vals: [num_rays, num_samples along ray]. Integration time.
rays_d: [num_rays, 3]. Direction of each ray.
Returns:
rgb_map: [num_rays, 3]. Estimated RGB color of a ray.
disp_map: [num_rays]. Disparity map. Inverse of depth map.
acc_map: [num_rays]. Sum of weights along each ray.
weights: [num_rays, num_samples]. Weights assigned to each sampled color.
depth_map: [num_rays]. Estimated distance to object.
"""
def raw2alpha(raw, dists, act_fn=F.relu): return 1. - \
torch.exp(-(act_fn(raw)+1e-6)*dists)
dists = z_vals[..., 1:] - z_vals[..., :-1]
dists = torch.cat([dists, torch.Tensor([1e10]).expand(
dists[..., :1].shape)], -1) # [N_rays, N_samples]
dists = dists * torch.norm(rays_d[..., None, :], dim=-1)
rgb = torch.sigmoid(raw[..., :3]) # [N_rays, N_samples, 3]
rgb = torch.cat((rgb[:, :-1, :], bc_rgb.unsqueeze(1)), dim=1)
noise = 0.
if raw_noise_std > 0.:
noise = torch.randn(raw[..., 3].shape) * raw_noise_std
# Overwrite randomly sampled data if pytest
if pytest:
np.random.seed(0)
noise = np.random.rand(*list(raw[..., 3].shape)) * raw_noise_std
noise = torch.Tensor(noise)
alpha = raw2alpha(raw[..., 3] + noise, dists) # [N_rays, N_samples]
# weights = alpha * tf.math.cumprod(1.-alpha + 1e-10, -1, exclusive=True)
weights = alpha * \
torch.cumprod(
torch.cat([torch.ones((alpha.shape[0], 1)), 1.-alpha + 1e-10], -1), -1)[:, :-1]
rgb_map = torch.sum(weights[..., None] * rgb, -2) # [N_rays, 3]
depth_map = torch.sum(weights * z_vals, -1)
disp_map = 1./torch.max(1e-10 * torch.ones_like(depth_map),
depth_map / torch.sum(weights, -1))
acc_map = torch.sum(weights, -1)
if white_bkgd:
rgb_map = rgb_map + (1.-acc_map[..., None])
return rgb_map, disp_map, acc_map, weights, depth_map
def render_rays(ray_batch,
bc_rgb,
aud_para,
network_fn,
network_query_fn,
N_samples,
retraw=False,
lindisp=False,
perturb=0.,
N_importance=0,
network_fine=None,
white_bkgd=False,
raw_noise_std=0.,
verbose=False,
pytest=False):
"""Volumetric rendering.
Args:
ray_batch: array of shape [batch_size, ...]. All information necessary
for sampling along a ray, including: ray origin, ray direction, min
dist, max dist, and unit-magnitude viewing direction.
network_fn: function. Model for predicting RGB and density at each point
in space.
network_query_fn: function used for passing queries to network_fn.
N_samples: int. Number of different times to sample along each ray.
retraw: bool. If True, include model's raw, unprocessed predictions.
lindisp: bool. If True, sample linearly in inverse depth rather than in depth.
perturb: float, 0 or 1. If non-zero, each ray is sampled at stratified
random points in time.
N_importance: int. Number of additional times to sample along each ray.
These samples are only passed to network_fine.
network_fine: "fine" network with same spec as network_fn.
white_bkgd: bool. If True, assume a white background.
raw_noise_std: ...
verbose: bool. If True, print more debugging info.
Returns:
rgb_map: [num_rays, 3]. Estimated RGB color of a ray. Comes from fine model.
disp_map: [num_rays]. Disparity map. 1 / depth.
acc_map: [num_rays]. Accumulated opacity along each ray. Comes from fine model.
raw: [num_rays, num_samples, 4]. Raw predictions from model.
rgb0: See rgb_map. Output for coarse model.
disp0: See disp_map. Output for coarse model.
acc0: See acc_map. Output for coarse model.
z_std: [num_rays]. Standard deviation of distances along ray for each
sample.
"""
N_rays = ray_batch.shape[0]
rays_o, rays_d = ray_batch[:, 0:3], ray_batch[:, 3:6] # [N_rays, 3] each
viewdirs = ray_batch[:, -3:] if ray_batch.shape[-1] > 8 else None
bounds = torch.reshape(ray_batch[..., 6:8], [-1, 1, 2])
near, far = bounds[..., 0], bounds[..., 1] # [-1,1]
t_vals = torch.linspace(0., 1., steps=N_samples)
if not lindisp:
z_vals = near * (1.-t_vals) + far * (t_vals)
else:
z_vals = 1./(1./near * (1.-t_vals) + 1./far * (t_vals))
z_vals = z_vals.expand([N_rays, N_samples])
if perturb > 0.:
# get intervals between samples
mids = .5 * (z_vals[..., 1:] + z_vals[..., :-1])
upper = torch.cat([mids, z_vals[..., -1:]], -1)
lower = torch.cat([z_vals[..., :1], mids], -1)
# stratified samples in those intervals
t_rand = torch.rand(z_vals.shape)
# Pytest, overwrite u with numpy's fixed random numbers
if pytest:
np.random.seed(0)
t_rand = np.random.rand(*list(z_vals.shape))
t_rand = torch.Tensor(t_rand)
t_rand[..., -1] = 1.0
z_vals = lower + (upper - lower) * t_rand
pts = rays_o[..., None, :] + rays_d[..., None, :] * \
z_vals[..., :, None] # [N_rays, N_samples, 3]
raw = network_query_fn(pts, viewdirs, aud_para, network_fn)
rgb_map, disp_map, acc_map, weights, depth_map = raw2outputs(
raw, z_vals, rays_d, bc_rgb, raw_noise_std, white_bkgd, pytest=pytest)
if N_importance > 0:
rgb_map_0, disp_map_0, acc_map_0 = rgb_map, disp_map, acc_map
z_vals_mid = .5 * (z_vals[..., 1:] + z_vals[..., :-1])
z_samples = sample_pdf(
z_vals_mid, weights[..., 1:-1], N_importance, det=(perturb == 0.), pytest=pytest)
z_samples = z_samples.detach()
z_vals, _ = torch.sort(torch.cat([z_vals, z_samples], -1), -1)
pts = rays_o[..., None, :] + rays_d[..., None, :] * \
z_vals[..., :, None] # [N_rays, N_samples + N_importance, 3]
run_fn = network_fn if network_fine is None else network_fine
raw = network_query_fn(pts, viewdirs, aud_para, run_fn)
rgb_map, disp_map, acc_map, weights, depth_map = raw2outputs(
raw, z_vals, rays_d, bc_rgb, raw_noise_std, white_bkgd, pytest=pytest)
ret = {'rgb_map': rgb_map, 'disp_map': disp_map, 'acc_map': acc_map}
if retraw:
ret['raw'] = raw
if N_importance > 0:
ret['rgb0'] = rgb_map_0
ret['disp0'] = disp_map_0
ret['acc0'] = acc_map_0
ret['z_std'] = torch.std(z_samples, dim=-1, unbiased=False) # [N_rays]
ret['last_weight'] = weights[..., -1]
for k in ret:
if (torch.isnan(ret[k]).any() or torch.isinf(ret[k]).any()) and DEBUG:
print(f"! [Numerical Error] {k} contains nan or inf.")
return ret
def config_parser():
import configargparse
parser = configargparse.ArgumentParser()
parser.add_argument('--config', is_config_file=True,
help='config file path')
parser.add_argument("--expname", type=str,
help='experiment name')
parser.add_argument("--basedir", type=str, default='./logs/',
help='where to store ckpts and logs')
parser.add_argument("--datadir", type=str, default='./data/llff/fern',
help='input data directory')
# training options
parser.add_argument("--netdepth", type=int, default=8,
help='layers in network')
parser.add_argument("--netwidth", type=int, default=256,
help='channels per layer')
parser.add_argument("--netdepth_fine", type=int, default=8,
help='layers in fine network')
parser.add_argument("--netwidth_fine", type=int, default=256,
help='channels per layer in fine network')
parser.add_argument("--N_rand", type=int, default=1024,
help='batch size (number of random rays per gradient step)')
parser.add_argument("--lrate", type=float, default=5e-4,
help='learning rate')
parser.add_argument("--lrate_decay", type=int, default=250,
help='exponential learning rate decay (in 1000 steps)')
parser.add_argument("--chunk", type=int, default=1024,
help='number of rays processed in parallel, decrease if running out of memory')
parser.add_argument("--netchunk", type=int, default=1024*64,
help='number of pts sent through network in parallel, decrease if running out of memory')
parser.add_argument("--no_batching", action='store_false',
help='only take random rays from 1 image at a time')
parser.add_argument("--no_reload", action='store_true',
help='do not reload weights from saved ckpt')
parser.add_argument("--ft_path", type=str, default=None,
help='specific weights npy file to reload for coarse network')
parser.add_argument("--N_iters", type=int, default=400000,
help='number of iterations')
# rendering options
parser.add_argument("--N_samples", type=int, default=16,
help='number of coarse samples per ray')
parser.add_argument("--N_importance", type=int, default=128,
help='number of additional fine samples per ray')
parser.add_argument("--perturb", type=float, default=1.,
help='set to 0. for no jitter, 1. for jitter')
parser.add_argument("--use_viewdirs", action='store_false',
help='use full 5D input instead of 3D')
parser.add_argument("--i_embed", type=int, default=0,
help='set 0 for default positional encoding, -1 for none')
parser.add_argument("--multires", type=int, default=10,
help='log2 of max freq for positional encoding (3D location)')
parser.add_argument("--multires_views", type=int, default=4,
help='log2 of max freq for positional encoding (2D direction)')
parser.add_argument("--raw_noise_std", type=float, default=0.,
help='std dev of noise added to regularize sigma_a output, 1e0 recommended')
# 只渲染,不进行优化等训练操作,相当于前向传播
parser.add_argument("--render_only", action='store_true',
help='do not optimize, reload weights and render out render_poses path')
parser.add_argument("--render_test", action='store_true',
help='render the test set instead of render_poses path')
parser.add_argument("--render_factor", type=int, default=0,
help='downsampling factor to speed up rendering, set 4 or 8 for fast preview')
# training options
parser.add_argument("--precrop_iters", type=int, default=0,
help='number of steps to train on central crops')
parser.add_argument("--precrop_frac", type=float,
default=.5, help='fraction of img taken for central crops')
# dataset options
parser.add_argument("--dataset_type", type=str, default='audface',
help='options: llff / blender / deepvoxels')
parser.add_argument("--testskip", type=int, default=8,
help='will load 1/N images from test/val sets, useful for large datasets like deepvoxels')
# deepvoxels flags
parser.add_argument("--shape", type=str, default='greek',
help='options : armchair / cube / greek / vase')
# blender flags
parser.add_argument("--white_bkgd", action='store_false',
help='set to render synthetic data on a white bkgd (always use for dvoxels)')
parser.add_argument("--half_res", action='store_true',
help='load blender synthetic data at 400x400 instead of 800x800')
# face flags
parser.add_argument("--with_test", type=int, default=0,
help='whether to use test set')
parser.add_argument("--dim_aud", type=int, default=64,
help='dimension of audio features for NeRF')
parser.add_argument("--sample_rate", type=float, default=0.95,
help="sample rate in a bounding box")
parser.add_argument("--near", type=float, default=0.3,
help="near sampling plane")
parser.add_argument("--far", type=float, default=0.9,
help="far sampling plane")
parser.add_argument("--test_file", type=str, default='transforms_test.json',
help='test file')
parser.add_argument("--aud_file", type=str, default='aud.npy',
help='test audio deepspeech file')
parser.add_argument("--win_size", type=int, default=16,
help="windows size of audio feature")
parser.add_argument("--smo_size", type=int, default=8,
help="window size for smoothing audio features")
parser.add_argument('--nosmo_iters', type=int, default=200000,
help='number of iterations befor applying smoothing on audio features')
# llff flags
parser.add_argument("--factor", type=int, default=8,
help='downsample factor for LLFF images')
parser.add_argument("--no_ndc", action='store_true',
help='do not use normalized device coordinates (set for non-forward facing scenes)')
parser.add_argument("--lindisp", action='store_true',
help='sampling linearly in disparity rather than depth')
parser.add_argument("--spherify", action='store_true',
help='set for spherical 360 scenes')
parser.add_argument("--llffhold", type=int, default=8,
help='will take every 1/N images as LLFF test set, paper uses 8')
# logging/saving options
parser.add_argument("--i_print", type=int, default=100,
help='frequency of console printout and metric loggin')
parser.add_argument("--i_img", type=int, default=500,
help='frequency of tensorboard image logging')
parser.add_argument("--i_weights", type=int, default=10000,
help='frequency of weight ckpt saving')
parser.add_argument("--i_testset", type=int, default=10000,
help='frequency of testset saving')
parser.add_argument("--i_video", type=int, default=50000,
help='frequency of render_poses video saving')
return parser
def train():
parser = config_parser()
args = parser.parse_args()
# Load data
'''
load_audface_data():
这个函数的目的是加载音频-视觉人脸识别任务的数据集。函数根据输入的参数和文件路径,加载训练集、验证集或测试集的图像数据、姿态数据、音频数据和其他相关信息。
具体来说,函数的目的是:
如果指定了测试集文件(test_file),则加载该测试集的元数据和音频特征数据。
函数返回测试集的姿态数据、音频数据、背景图像、图像尺寸和中心点坐标等信息。
如果未指定测试集文件,函数会加载训练集和验证集的元数据和音频特征数据。函数将训练集和验证集的图像数据、姿态数据、音频数据和样本矩形区域数据合并,
并返回合并后的数据以及背景图像、图像尺寸和中心点坐标等信息。
因此,该函数的主要目的是加载用于音频-视觉人脸识别任务的数据集,并提供方便访问和使用数据的接口。
'''
args.datadir = '/home/sza/AD-NeRF/dataset/Obama' # config.txt文件
args.basedir = '/home/sza/AD-NeRF/dataset/Obama/logs' # 指定实验结果的输出路径
if args.dataset_type == 'audface':
if args.with_test == 1:# 是否使用测试集
poses, auds, bc_img, hwfcxy = \
load_audface_data(args.datadir, args.testskip,
args.test_file, args.aud_file)
images = np.zeros(1)
else:
images, poses, auds, bc_img, hwfcxy, sample_rects, mouth_rects, i_split = load_audface_data(
args.datadir, args.testskip)
print('Loaded audface', images.shape, hwfcxy, args.datadir)
if args.with_test == 0:
i_train, i_val = i_split
near = args.near
far = args.far
else:
print('Unknown dataset type', args.dataset_type, 'exiting')
return
# Cast intrinsics to right types
# 将内在函数转换为正确的类型
H, W, focal, cx, cy = hwfcxy
H, W = int(H), int(W)
hwf = [H, W, focal]
hwfcxy = [H, W, focal, cx, cy]
# if args.render_test:
# render_poses = np.array(poses[i_test])
# Create log dir and copy the config file
'''
创建日志目录并复制配置文件。
首先,代码获取了args.basedir和args.expname的值,用于构建日志目录的路径。args.basedir是指定实验结果的输出路径,args.expname是实验名称。
'''
args.expname = 'Obama_com'
basedir = args.basedir
absolute_path = os.path.abspath(basedir)
expname = args.expname
'''
使用os.makedirs函数创建日志目录,如果该目录已经存在,则不进行任何操作。
这里在expname目录下生成args.txt和config.txt文件
其中args.txt文件是是用来保存之前设置的所有arg配置
for arg in sorted(vars(args)):
attr = getattr(args, arg)
file.write('{} = {}\n'.format(arg, attr))
而config则是读取制作数据集的产生的文件中'
HeadNeRF_config.txt
TorsoNeRF_config.txt
TorsoNeRFTest_config.txt
'''
os.makedirs(os.path.join(basedir, expname), exist_ok=True)
f = os.path.join(basedir, expname, 'args.txt')
with open(f, 'w') as file:
for arg in sorted(vars(args)):
attr = getattr(args, arg)
file.write('{} = {}\n'.format(arg, attr))
if args.config is not None:
f = os.path.join(basedir, expname, 'config.txt')
with open(f, 'w') as file:
file.write(open(args.config, 'r').read())
# Create nerf model
# 通过这段代码,NERF模型、音频网络(`AudNet`)和音频注意力网络(`AudAttNet`)被创建,并加载了之前训练的状态,以便从之前的训练中继续训练或进行推断。
'''
首先,调用`create_nerf(args)`函数创建NERF模型所需的参数和对象。
该函数返回了训练和测试的渲染参数`render_kwargs_train`和`render_kwargs_test`,
优化器`optimizer`,起始训练步骤`start`,可训练变量`grad_vars`,学习到的代码`learned_codes`,以及与音频网络(`AudNet`)相关的状态和优化器状态。
'''
render_kwargs_train, render_kwargs_test, start, grad_vars, optimizer, \
learned_codes, AudNet_state, optimizer_aud_state, AudAttNet_state, optimizer_audatt_state \
= create_nerf(args)
global_step = start
'''
AudioAttNet模型具有注意力机制,而AudioNet模型是一个编码器模型。
根据给定的参数创建音频网络(`AudNet`)和音频注意力网络(`AudAttNet`)。
`AudNet`是一个接收音频输入并生成相应特征的网络模型,而`AudAttNet`是一个用于音频注意力的网络模型。这两个网络被移动到指定的设备(如GPU)上。
'''
AudNet = AudioNet(args.dim_aud, args.win_size).to(device)
AudAttNet = AudioAttNet().to(device)
'''
使用`torch.optim.Adam`函数创建两个优化器`optimizer_Aud`和`optimizer_AudAtt`,用于分别优化`AudNet`和`AudAttNet`的参数。这些优化器将被用于在训练过程中更新网络的权重。
'''
optimizer_Aud = torch.optim.Adam(
params=list(AudNet.parameters()), lr=args.lrate, betas=(0.9, 0.999))
optimizer_AudAtt = torch.optim.Adam(
params=list(AudAttNet.parameters()), lr=args.lrate, betas=(0.9, 0.999))
'''
加载先前保存的模型状态字典(`AudNet_state`、`optimizer_aud_state`、`AudAttNet_state`和`optimizer_audatt_state`),
将之前训练的模型状态恢复到相应的网络和优化器中。如果之前没有保存的模型状态,则跳过此步骤。
'''
if AudNet_state is not None:
AudNet.load_state_dict(AudNet_state, strict=False)
if optimizer_aud_state is not None:
optimizer_Aud.load_state_dict(optimizer_aud_state)
if AudAttNet_state is not None:
AudAttNet.load_state_dict(AudAttNet_state, strict=False)
if optimizer_audatt_state is not None:
optimizer_AudAtt.load_state_dict(optimizer_audatt_state)
bds_dict = {
'near': near,
'far': far,
}
'''
将近场和远场的边界距离(`near`和`far`)作为字典添加到`render_kwargs_train`和`render_kwargs_test`中。这些参数将用于指定NERF模型的工作范围。
'''
render_kwargs_train.update(bds_dict)
render_kwargs_test.update(bds_dict)
# Move training data to GPU
bc_img = torch.Tensor(bc_img).to(device).float()/255.0
poses = torch.Tensor(poses).to(device).float()
auds = torch.Tensor(auds).to(device).float()
# 只渲染,不进行优化等训练操作,相当于前向传播
if args.render_only:# 默认为False
print('RENDER ONLY')
with torch.no_grad():
# Default is smoother render_poses path
images = None
testsavedir = os.path.join(basedir, expname, 'renderonly_{}_{:06d}'.format(
'test' if args.render_test else 'path', start))
os.makedirs(testsavedir, exist_ok=True)
print('test poses shape', poses.shape)
auds_val = AudNet(auds)
rgbs, disp, last_weight = render_path(poses, auds_val, bc_img, hwfcxy, args.chunk, render_kwargs_test,
gt_imgs=images, savedir=testsavedir, render_factor=args.render_factor)
np.save(os.path.join(testsavedir, 'last_weight.npy'), last_weight)
print('Done rendering', testsavedir)
imageio.mimwrite(os.path.join(
testsavedir, 'video.mp4'), to8b(rgbs), fps=30, quality=8)
return
num_frames = images.shape[0]
'''
通过这段代码,NERF模型、音频网络(`AudNet`)和音频注意力网络(`AudAttNet`)被创建,并加载了之前训练的状态,以便从之前的训练中继续训练或进行推断。
'''
# Prepare raybatch tensor if batching random rays
N_rand = args.N_rand
print('N_rand', N_rand, 'no_batching',
args.no_batching, 'sample_rate', args.sample_rate)
use_batching = not args.no_batching
if use_batching:
# For random ray batching
print('get rays')
rays = np.stack([get_rays_np(H, W, focal, p, cx, cy)
for p in poses[:, :3, :4]], 0) # [N, ro+rd, H, W, 3]
print('done, concats')
# [N, ro+rd+rgb, H, W, 3]
rays_rgb = np.concatenate([rays, images[:, None]], 1)
# [N, H, W, ro+rd+rgb, 3]
rays_rgb = np.transpose(rays_rgb, [0, 2, 3, 1, 4])
rays_rgb = np.stack([rays_rgb[i]
for i in i_train], 0) # train images only
# [(N-1)*H*W, ro+rd+rgb, 3]
rays_rgb = np.reshape(rays_rgb, [-1, 3, 3])
rays_rgb = rays_rgb.astype(np.float32)
print('shuffle rays')
np.random.shuffle(rays_rgb)
print('done')
i_batch = 0
if use_batching:
rays_rgb = torch.Tensor(rays_rgb).to(device)
N_iters = args.N_iters + 1
print('Begin')
print('TRAIN views are', i_train)
print('VAL views are', i_val)
start = start + 1
for i in trange(start, N_iters):
time0 = time.time()
# Sample random ray batch
if use_batching:
# Random over all images
batch = rays_rgb[i_batch:i_batch+N_rand] # [B, 2+1, 3*?]
batch = torch.transpose(batch, 0, 1)
batch_rays, target_s = batch[:2], batch[2]
i_batch += N_rand
if i_batch >= rays_rgb.shape[0]:
print("Shuffle data after an epoch!")
rand_idx = torch.randperm(rays_rgb.shape[0])
rays_rgb = rays_rgb[rand_idx]
i_batch = 0
else:
'''
从i_train数组中随机选择一个索引img_i,表示选择一个训练图像。
'''
# Random from one image
img_i = np.random.choice(i_train)
'''
使用imageio.imread()函数读取图像数据,并将其转换为torch.Tensor类型,并将像素值缩放到范围[0, 1]之间。
'''
target = torch.as_tensor(imageio.imread(
images[img_i])).to(device).float()/255.0
'''
从poses数组中获取对应的姿态矩阵pose。
从sample_rects数组中获取采样矩形区域rect。
从mouth_rects数组中获取嘴部矩形区域mouth_rect。
从auds数组中获取对应的音频特征aud。
'''
pose = poses[img_i, :3, :4]
rect = sample_rects[img_i]
mouth_rect = mouth_rects[img_i]
aud = auds[img_i]
if global_step >= args.nosmo_iters:# 我这里是站global=149999 nosmo_iters=200000
'''
如果global_step大于等于args.nosmo_iters,进行平滑处理:
根据args.smo_size的一半大小获取左右相邻的音频特征窗口范围。
对超出范围的部分进行填充。
使用AudAttNet对音频特征窗口进行处理,得到平滑后的音频特征aud_smo。
'''
smo_half_win = int(args.smo_size / 2)
left_i = img_i - smo_half_win
right_i = img_i + smo_half_win
pad_left, pad_right = 0, 0
if left_i < 0:
pad_left = -left_i
left_i = 0
if right_i > i_train.shape[0]:
pad_right = right_i-i_train.shape[0]
right_i = i_train.shape[0]
auds_win = auds[left_i:right_i]
if pad_left > 0:
auds_win = torch.cat(
(torch.zeros_like(auds_win)[:pad_left], auds_win), dim=0)
if pad_right > 0:
auds_win = torch.cat(
(auds_win, torch.zeros_like(auds_win)[:pad_right]), dim=0)
auds_win = AudNet(auds_win)
aud = auds_win[smo_half_win]
aud_smo = AudAttNet(auds_win)
else:
'''
如果global_step小于args.nosmo_iters,直接使用AudNet对音频特征进行处理,得到处理后的音频特征aud。
'''
aud = AudNet(aud.unsqueeze(0))
if N_rand is not None:
'''
如果N_rand不为None,进行随机采样:
根据图像尺寸、相机参数和姿态矩阵获取射线的起点和方向。
如果迭代次数小于args.precrop_iters,进行中心裁剪,裁剪尺寸为图像尺寸的一半乘以args.precrop_frac。
对于每个采样点,判断是否在rect区域内,将其分为在区域内和在区域外的点。
根据采样率选择在区域内和在区域外的点,使得采样点数满足N_rand要求。
根据选择的采样点坐标从目标图像和背景图像中提取对应的像素值。
'''
rays_o, rays_d = get_rays(
H, W, focal, torch.Tensor(pose), cx, cy) # (H, W, 3), (H, W, 3)
if i < args.precrop_iters:
dH = int(H//2 * args.precrop_frac)
dW = int(W//2 * args.precrop_frac)
coords = torch.stack(
torch.meshgrid(
torch.linspace(H//2 - dH, H//2 + dH - 1, 2*dH),
torch.linspace(W//2 - dW, W//2 + dW - 1, 2*dW)
), -1)
if i == start:
print(
f"[Config] Center cropping of size {2*dH} x {2*dW} is enabled until iter {args.precrop_iters}")
else:
coords = torch.stack(torch.meshgrid(torch.linspace(
0, H-1, H), torch.linspace(0, W-1, W)), -1) # (H, W, 2)
coords = torch.reshape(coords, [-1, 2]) # (H * W, 2)
if args.sample_rate > 0:
rect_inds = (coords[:, 0] >= rect[0]) & (
coords[:, 0] <= rect[0] + rect[2]) & (
coords[:, 1] >= rect[1]) & (
coords[:, 1] <= rect[1] + rect[3])
coords_rect = coords[rect_inds]
coords_norect = coords[~rect_inds]
rect_num = int(N_rand*args.sample_rate)
norect_num = N_rand - rect_num
select_inds_rect = np.random.choice(
coords_rect.shape[0], size=[rect_num], replace=False) # (N_rand,)
# (N_rand, 2)
select_coords_rect = coords_rect[select_inds_rect].long()
select_inds_norect = np.random.choice(
coords_norect.shape[0], size=[norect_num], replace=False) # (N_rand,)
# (N_rand, 2)
select_coords_norect = coords_norect[select_inds_norect].long(
)
select_coords = torch.cat(
(select_coords_rect, select_coords_norect), dim=0)
else:
select_inds = np.random.choice(
coords.shape[0], size=[N_rand], replace=False) # (N_rand,)
select_coords = coords[select_inds].long()
'''
返回采样得到的数据,包括射线起点、射线方向、目标图像的采样像素值和背景图像的采样像素值。
'''
rays_o = rays_o[select_coords[:, 0],
select_coords[:, 1]] # (N_rand, 3)
rays_d = rays_d[select_coords[:, 0],
select_coords[:, 1]] # (N_rand, 3)
batch_rays = torch.stack([rays_o, rays_d], 0)
target_s = target[select_coords[:, 0],
select_coords[:, 1]] # (N_rand, 3)
bc_rgb = bc_img[select_coords[:, 0],
select_coords[:, 1]]
##### Core optimization loop #####
'''
这段代码是核心的优化循环,用于执行模型的前向传播和计算损失。
'''
if global_step >= args.nosmo_iters:
'''
如果global_step大于等于args.nosmo_iters,使用平滑后的音频特征aud_smo进行渲染;
'''
rgb, disp, acc, _, extras = render_dynamic_face(H, W, focal, cx, cy, chunk=args.chunk, rays=batch_rays,
aud_para=aud_smo, bc_rgb=bc_rgb,
verbose=i < 10, retraw=True,
**render_kwargs_train)
else:
'''
否则,使用原始音频特征aud进行渲染。调用render_dynamic_face()函数进行面部渲染,根据给定的参数生成渲染的RGB图像、视差图、累积图以及其他附加信息。
'''
rgb, disp, acc, _, extras = render_dynamic_face(H, W, focal, cx, cy, chunk=args.chunk, rays=batch_rays,
aud_para=aud, bc_rgb=bc_rgb,
verbose=i < 10, retraw=True,
**render_kwargs_train)
'''
使用优化器对象optimizer、optimizer_Aud和optimizer_AudAtt将梯度置零,准备进行反向传播和参数更新。
'''
optimizer.zero_grad()
optimizer_Aud.zero_grad()
optimizer_AudAtt.zero_grad()
'''
计算图像损失,使用img2mse()函数计算渲染的RGB图像与目标图像之间的均方误差(MSE)损失。
'''
img_loss = img2mse(rgb, target_s)
trans = extras['raw'][..., -1]
loss = img_loss
'''
使用mse2psnr()函数计算图像损失的峰值信噪比(PSNR)。
'''
psnr = mse2psnr(img_loss)
if 'rgb0' in extras:
img_loss0 = img2mse(extras['rgb0'], target_s)
loss = loss + img_loss0
psnr0 = mse2psnr(img_loss0)
'''
使用loss.backward()进行损失的反向传播,计算梯度。
'''
loss.backward()
'''
使用optimizer.step()和optimizer_Aud.step()分别更新nerf模型和音频特征模型的参数。
'''
optimizer.step()
optimizer_Aud.step()
if global_step >= args.nosmo_iters:
'''
如果global_step大于等于args.nosmo_iters,使用optimizer_AudAtt.step()更新注意力模型的参数。
'''
optimizer_AudAtt.step()
# NOTE: IMPORTANT!
### update learning rate ###
''''
定义学习率衰减的参数。decay_rate表示衰减率,取值为0.1,表示每个衰减步骤学习率减小到原来的10%。
decay_steps表示衰减步骤的总数,计算为args.lrate_decay乘以1000,
表示每1000个步骤进行一次学习率衰减。
'''
decay_rate = 0.1
decay_steps = args.lrate_decay * 1000
'''
根据当前的全局步骤global_step和衰减步骤数,计算新的学习率new_lrate。采用指数衰减的方式,将初始学习率args.lrate乘以衰减率的指数函数。
'''
new_lrate = args.lrate * (decay_rate ** (global_step / decay_steps))
for param_group in optimizer.param_groups:
'''
遍历优化器optimizer的参数组,并将学习率更新为new_lrate。
'''
param_group['lr'] = new_lrate
for param_group in optimizer_Aud.param_groups:
'''
遍历音频特征优化器optimizer_Aud的参数组,并将学习率更新为new_lrate。
'''
param_group['lr'] = new_lrate
for param_group in optimizer_AudAtt.param_groups:
'''
遍历注意力模型优化器optimizer_AudAtt的参数组,并将学习率更新为new_lrate的5倍。
'''
param_group['lr'] = new_lrate*5
################################
dt = time.time()-time0
# Rest is logging
if i % args.i_weights == 0:
path = os.path.join(basedir, expname, '{:06d}_head.tar'.format(i))
torch.save({
'global_step': global_step,
'network_fn_state_dict': render_kwargs_train['network_fn'].state_dict(),
'network_fine_state_dict': render_kwargs_train['network_fine'].state_dict(),
'network_audnet_state_dict': AudNet.state_dict(),
'optimizer_state_dict': optimizer.state_dict(),
'optimizer_aud_state_dict': optimizer_Aud.state_dict(),
'network_audattnet_state_dict': AudAttNet.state_dict(),
'optimizer_audatt_state_dict': optimizer_AudAtt.state_dict(),
}, path)
print('Saved checkpoints at', path)
if i % args.i_testset == 0 and i > 0:
testsavedir = os.path.join(
basedir, expname, 'testset_{:06d}'.format(i))
os.makedirs(testsavedir, exist_ok=True)
print('test poses shape', poses[i_val].shape)
auds_val = AudNet(auds[i_val])
with torch.no_grad():
render_path(torch.Tensor(poses[i_val]).to(
device), auds_val, bc_img, hwfcxy, args.chunk, render_kwargs_test, gt_imgs=None, savedir=testsavedir)
print('Saved test set')
if i % args.i_print == 0:
tqdm.write(
f"[TRAIN] Iter: {i} Loss: {loss.item()} PSNR: {psnr.item()}")
global_step += 1
if __name__ == '__main__':
torch.set_default_tensor_type('torch.cuda.FloatTensor')
train()
到了这里,关于[AD-NeRF]由音频和人脸图像合成人脸视频并表现出自然的说话风格------30系显卡实现的文章就介绍完了。如果您还想了解更多内容,请在右上角搜索TOY模板网以前的文章或继续浏览下面的相关文章,希望大家以后多多支持TOY模板网!