"""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)

"""Prepares inputs and applies network 'fn'.

inputs_flat = torch.reshape(inputs, [-1, inputs.shape[-1]])

embedded = embed_fn(inputs_flat)

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]])

"""Render rays in smaller minibatches to avoid OOM.

all_ret = {}

for i in range(0, rays_flat.shape[0], chunk):

#渲染光线，得到rgb_map等封装成的字典

ret = render_rays(rays_flat[i:i+chunk], **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}

near=0., far=1.,

use_viewdirs=False, c2w_staticcam=None,

**kwargs):

"""Render rays

if c2w is not None:

# special case to render full image

#渲染图像中的所有点

rays_o, rays_d = get_rays(H, W, K, c2w)

else:

# use provided ray batch

rays_o, rays_d = rays

if use_viewdirs:

# provide ray directions as input

#使用视角信息，将rays_d正则化赋值给viewdirs

viewdirs = rays_d

if c2w_staticcam is not None:

# special case to visualize effect of viewdirs

rays_o, rays_d = get_rays(H, W, K, c2w_staticcam)

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, K[0][0], 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])

#（1024,8）8：3+3+1+1

rays = torch.cat([rays_o, rays_d, near, far], -1)

if use_viewdirs:

#如果使用视角就加上视角

rays = torch.cat([rays, viewdirs], -1)

# Render and reshape

#开始并行计算光线属性，返回一堆rgb_map等

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}

H, W, focal = hwf

if render_factor!=0:

# Render downsampled for speed

H = H//render_factor

W = W//render_factor

focal = focal/render_factor

rgbs = []

disps = []

t = time.time()

for i, c2w in enumerate(tqdm(render_poses)):

print(i, time.time() - t)

t = time.time()

rgb, disp, acc, _ = render(H, W, K, chunk=chunk, c2w=c2w[:3,:4], **render_kwargs)

rgbs.append(rgb.cpu().numpy())

disps.append(disp.cpu().numpy())

if i==0:

print(rgb.shape, disp.shape)

"""

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)

"""Instantiate NeRF's MLP model.

'''

#对3D位置坐标编码

embed_fn, input_ch = get_embedder(args.multires, args.i_embed)

input_ch_views = 0

embeddirs_fn = None

if args.use_viewdirs:

#对视角2D坐标编码

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，输出为RGBA

#use_viewdirs表示使用视角信息（即使用5D而非3D坐标）

model = NeRF(D=args.netdepth, W=args.netwidth,

input_ch=input_ch, 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

if args.N_importance > 0:

model_fine = NeRF(D=args.netdepth_fine, W=args.netwidth_fine,

input_ch=input_ch, 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 = lambda inputs, viewdirs, network_fn : run_network(inputs, viewdirs, network_fn,

embed_fn=embed_fn,

embeddirs_fn=embeddirs_fn,

netchunk=args.netchunk)

# 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

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 sorted(os.listdir(os.path.join(basedir, expname))) if 'tar' in f]

print('Found ckpts', ckpts)

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'])

# 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'])

##########################

#初始化渲染的一些参数

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, #为True则透明部分变成白色

'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.

"""Transforms model's predictions to semantically meaningful values.

raw2alpha = lambda raw, dists, act_fn=F.relu: 1.-torch.exp(-act_fn(raw)*dists)

#计算相邻两点z轴之间的距离（1024,63）

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值

rgb = torch.sigmoid(raw[...,:3]) # [N_rays, N_samples, 3]

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)

#计算分配给每个采样点的权重，用来集合成为预测rgb （1024,64）

weights = alpha * torch.cumprod(torch.cat([torch.ones((alpha.shape[0], 1)), 1.-alpha + 1e-10], -1), -1)[:, :-1]

#每个ray有个预测的rgb值

rgb_map = torch.sum(weights[...,None] * rgb, -2) # [N_rays, 3]

#（1024）估计每个ray对应的点到物体的距离

depth_map = torch.sum(weights * z_vals, -1)

#视差 （1024）

disp_map = 1./torch.max(1e-10 * torch.ones_like(depth_map), depth_map / torch.sum(weights, -1))

#每条ray的权重加和 （1024）

acc_map = torch.sum(weights, -1)

if white_bkgd:

rgb_map = rgb_map + (1.-acc_map[...,None])

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.

#将各种参数分离出来

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]

#N_samples是每条rays上采样的次数，t_vals就是一个长为N_samples的向量

t_vals = torch.linspace(0., 1., steps=N_samples)

#确定z轴的值

if not lindisp:

#根据depth线性采样

#（1024,64） 1024束光线，每个随机采样64个点

z_vals = near * (1.-t_vals) + far * (t_vals)

else:

#按照depth的inverse线性采样

z_vals = 1./(1./near * (1.-t_vals) + 1./far * (t_vals))

#为什么又来一句丝毫没有用的代码？？？？

#expand是用来扩展维度的

z_vals = z_vals.expand([N_rays, N_samples])

if perturb > 0.:

#分层随机采样

# get intervals between samples

#（1024,63）相当于同一条的ray的前一个采样点与后一个采样点相加*0.5

mids = .5 * (z_vals[...,1:] + z_vals[...,:-1])

#（1024,64）与最后一个采样点拼接起来，确定上界

upper = torch.cat([mids, z_vals[...,-1:]], -1)

#（1024,64）与第一个采样点拼接起来，确定下界

lower = torch.cat([z_vals[...,:1], mids], -1)

# stratified samples in those intervals

#从（0,1）的均匀分布中抽取一组随机数，形状为z_vals.shape，即（1024，,6）

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相当于在每个区间内的比例，这里算出实际采样点的z轴的值

z_vals = lower + (upper - lower) * t_rand

#生成光线上每个采样点的位置：pts=o+z*d，所以z解释称绝对长度更合理

#（1024,64,3）1024个光线每条光线有64个采样点，每个采样点有（x,y,z）坐标

pts = rays_o[...,None,:] + rays_d[...,None,:] * z_vals[...,:,None] # [N_rays, N_samples, 3]

#raw = run_network(pts)

#将参数传给network_fn中前向传播得到每个点对应的RGBA

raw = network_query_fn(pts, viewdirs, network_fn)

#计算得到rgb、深度图等等

rgb_map, disp_map, acc_map, weights, depth_map = raw2outputs(raw, z_vals, rays_d, 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, run_fn)

rgb_map, disp_map, acc_map, weights, depth_map = raw2outputs(raw, z_vals, rays_d, 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]

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.")

#基本参数

import configargparse

parser = configargparse.ArgumentParser()

#生成config.txt文件

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')

#fine network和network的区别？

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=32*32*4,

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*32,

help='number of rays processed in parallel, decrease if running out of memory')

#what is pts？

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_true',

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')

# rendering options

#渲染参数

#每条射线的粗样本数

parser.add_argument("--N_samples", type=int, default=64,

help='number of coarse samples per ray')

#细样本和粗样本的区别？

#每条射线额外的细样本数

parser.add_argument("--N_importance", type=int, default=0,

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_true',

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')

#多分辨率(3D)

parser.add_argument("--multires", type=int, default=10,

help='log2 of max freq for positional encoding (3D location)')

#(2D)

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')

#crop占比

parser.add_argument("--precrop_frac", type=float,

default=.5, help='fraction of img taken for central crops')

# dataset options

#数据集选择 voxel：体素

parser.add_argument("--dataset_type", type=str, default='llff',

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_true',

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')

## 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=50000,

help='frequency of testset saving')

parser.add_argument("--i_video", type=int, default=50000,

help='frequency of render_poses video saving')

parser = config_parser()

args = parser.parse_args()

# Load data，根据传入的参数对数据集做各种处理

K = None

if args.dataset_type == 'llff':

#这些shape是从哪里看出来的？？？

#shape: images[20,378,504,3],poses[20,3,5],render_poses[120,3,5]

images, poses, bds, render_poses, i_test = load_llff_data(args.datadir, args.factor,

recenter=True, bd_factor=.75,

spherify=args.spherify)

#hwf相当于把poses最后一个通道的第一行的所有元素拿出来

hwf = poses[0,:3,-1]

#去掉最后一个通道

poses = poses[:,:3,:4]

print('Loaded llff', images.shape, render_poses.shape, hwf, args.datadir)

if not isinstance(i_test, list):

i_test = [i_test]

#分配数据集

if args.llffhold > 0:

print('Auto LLFF holdout,', args.llffhold)

i_test = np.arange(images.shape[0])[::args.llffhold]

#val=test

i_val = i_test

i_train = np.array([i for i in np.arange(int(images.shape[0])) if

(i not in i_test and i not in i_val)])

print('DEFINING BOUNDS')

#定义边界

if args.no_ndc:

near = np.ndarray.min(bds) * .9

far = np.ndarray.max(bds) * 1.

else:

near = 0.

far = 1.

print('NEAR FAR', near, far)

#blender

elif args.dataset_type == 'blender':

images, poses, render_poses, hwf, i_split = load_blender_data(args.datadir, args.half_res, args.testskip)

print('Loaded blender', images.shape, render_poses.shape, hwf, args.datadir)

i_train, i_val, i_test = i_split

#

near = 2.

far = 6.

if args.white_bkgd:

images = images[...,:3]*images[...,-1:] + (1.-images[...,-1:])

else:

images = images[...,:3]

#linemod

elif args.dataset_type == 'LINEMOD':

images, poses, render_poses, hwf, K, i_split, near, far = load_LINEMOD_data(args.datadir, args.half_res, args.testskip)

print(f'Loaded LINEMOD, images shape: {images.shape}, hwf: {hwf}, K: {K}')

print(f'[CHECK HERE] near: {near}, far: {far}.')

i_train, i_val, i_test = i_split

if args.white_bkgd:

images = images[...,:3]*images[...,-1:] + (1.-images[...,-1:])

else:

images = images[...,:3]

#deepvoxels

elif args.dataset_type == 'deepvoxels':

images, poses, render_poses, hwf, i_split = load_dv_data(scene=args.shape,

basedir=args.datadir,

testskip=args.testskip)

print('Loaded deepvoxels', images.shape, render_poses.shape, hwf, args.datadir)

i_train, i_val, i_test = i_split

hemi_R = np.mean(np.linalg.norm(poses[:,:3,-1], axis=-1))

near = hemi_R-1.

far = hemi_R+1.

else:

print('Unknown dataset type', args.dataset_type, 'exiting')

return

# Cast intrinsics to right types

H, W, focal = hwf

H, W = int(H), int(W)

hwf = [H, W, focal]

#拼接出内参矩阵K

if K is None:

K = np.array([

[focal, 0, 0.5*W],

[0, focal, 0.5*H],

[0, 0, 1]

])

#这个参数表示拿测试集的pose来做渲染

if args.render_test:

render_poses = np.array(poses[i_test])

# Create log dir and copy the config file

basedir = args.basedir

expname = args.expname

os.makedirs(os.path.join(basedir, expname), exist_ok=True)

# 拼接出参数指定的存储log的路径

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:

# config路径

f = os.path.join(basedir, expname, 'config.txt')

with open(f, 'w') as file:

file.write(open(args.config, 'r').read())

# Create nerf model

#生成网络

render_kwargs_train, render_kwargs_test, start, grad_vars, optimizer = create_nerf(args)

global_step = start

bds_dict = {

'near' : near,

'far' : far,

}

#加入边界

render_kwargs_train.update(bds_dict)

render_kwargs_test.update(bds_dict)

# Move testing data to GPU

render_poses = torch.Tensor(render_poses).to(device)

# Short circuit if only rendering out from trained model

#只渲染并生成视频

if args.render_only:

print('RENDER ONLY')

with torch.no_grad():

if args.render_test:

# render_test switches to test poses

images = images[i_test]

else:

# 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', render_poses.shape)

rgbs, _ = render_path(render_poses, hwf, K, args.chunk, render_kwargs_test, gt_imgs=images, savedir=testsavedir, render_factor=args.render_factor)

print('Done rendering', testsavedir)

imageio.mimwrite(os.path.join(testsavedir, 'video.mp4'), to8b(rgbs), fps=30, quality=8)

#直接返回，训练结束

return

# Prepare raybatch tensor if batching random rays

#N_rand表示每批有多大

N_rand = args.N_rand

use_batching = not args.no_batching

#批处理，从一批图像中采取光线

#为了看到下面的代码暂时改一下

#use_batching=True

if use_batching:

# For random ray batching

print('get rays')

#（138,2,400,400,3）2：前面是rays_o，后面是rays_d

rays = np.stack([get_rays_np(H, W, K, p) for p in poses[:,:3,:4]], 0) # [N, ro+rd, H, W, 3]

print('done, concats')

#（138,3,400,400,3）3：加了一个rgb，刚好也是（400,400,3）

#None是一个常用的给array/tensor添加维度的手段

rays_rgb = np.concatenate([rays, images[:,None]], 1) # [N, ro+rd+rgb, H, W, 3]

#调换一下位置 (138,400,400,3,3)

rays_rgb = np.transpose(rays_rgb, [0,2,3,1,4]) # [N, H, W, ro+rd+rgb, 3]

#只把train的挑出来，下面就都是只有train的了，只有100个

rays_rgb = np.stack([rays_rgb[i] for i in i_train], 0) # train images only

#（100*400*400,3,3）

rays_rgb = np.reshape(rays_rgb, [-1,3,3]) # [(N-1)*H*W, ro+rd+rgb, 3]

rays_rgb = rays_rgb.astype(np.float32)

print('shuffle rays')

np.random.shuffle(rays_rgb)

print('done')

i_batch = 0

# Move training data to GPU

if use_batching:

images = torch.Tensor(images).to(device)

poses = torch.Tensor(poses).to(device)

if use_batching:

rays_rgb = torch.Tensor(rays_rgb).to(device)

#默认迭代200000次

N_iters = 200000 + 1

print('Begin')

print('TRAIN views are', i_train)

print('TEST views are', i_test)

print('VAL views are', i_val)

# Summary writers

# writer = SummaryWriter(os.path.join(basedir, 'summaries', expname))

start = start + 1

#trange=tqdm(range()),可以打印进度条

for i in trange(start, N_iters):

#下面都是每个epoch做的事情

time0 = time.time()

# Sample random ray batch

if use_batching:

#分批加载rays

# Random over all images

batch = rays_rgb[i_batch:i_batch+N_rand] # [B, 2+1, 3]

batch = torch.transpose(batch, 0, 1) #[3,B,3] 第一列和第二列互换

batch_rays, target_s = batch[:2], batch[2] #把rgb和对应像素点颜色分离

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:

# Random from one image

#随机选择一张图像进行训练

img_i = np.random.choice(i_train) #随机选图片的index

target = images[img_i]

target = torch.Tensor(target).to(device)

pose = poses[img_i, :3,:4]

if N_rand is not None:

#和get_rays_up实现过程类似，不过是对单张图生成rays

rays_o, rays_d = get_rays(H, W, K, torch.Tensor(pose)) # (H, W, 3), (H, W, 3)

#precrop这两个参数都和中心裁剪有关

#计算笛卡尔坐标

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就是笛卡尔坐标

coords = torch.reshape(coords, [-1,2]) # (H * W, 2)（x，y）

#选一批光线，也是从400*400个点中选取N_rand个点

select_inds = np.random.choice(coords.shape[0], size=[N_rand], replace=False) # (N_rand,)

#取出这些点的坐标

select_coords = coords[select_inds].long() # (N_rand, 2)

#select_coords[:,0]是横坐标，另一个是纵坐标，这样可以取出这些点的rays_o

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)

##### Core optimization loop #####

__import__('ipdb').set_trace()

#得到生成的光线后开始对光线进行渲染工作

rgb, disp, acc, extras = render(H, W, K, chunk=args.chunk, rays=batch_rays,

verbose=i < 10, retraw=True,

**render_kwargs_train)

optimizer.zero_grad()

img_loss = img2mse(rgb, target_s) #与target图像计算MSE损失

trans = extras['raw'][...,-1]

loss = img_loss #所以loss是预测的rgb与实际rgb的MSE

psnr = mse2psnr(img_loss) #将MSE损失转化为PSNR

if 'rgb0' in extras:

img_loss0 = img2mse(extras['rgb0'], target_s)

loss = loss + img_loss0

psnr0 = mse2psnr(img_loss0)

loss.backward()

optimizer.step()

# NOTE: IMPORTANT!

### update learning rate ###

'''

decay_rate = 0.1 #衰减率

decay_steps = args.lrate_decay * 1000

new_lrate = args.lrate * (decay_rate ** (global_step / decay_steps))

for param_group in optimizer.param_groups:

param_group['lr'] = new_lrate

################################

dt = time.time()-time0

# print(f"Step: {global_step}, Loss: {loss}, Time: {dt}")

##### end #####

#不看logging，嘻嘻

# Rest is logging

if i%args.i_weights==0:

path = os.path.join(basedir, expname, '{:06d}.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(),

'optimizer_state_dict': optimizer.state_dict(),

}, path)

print('Saved checkpoints at', path)

if i%args.i_video==0 and i > 0:

# Turn on testing mode

with torch.no_grad():

rgbs, disps = render_path(render_poses, hwf, K, args.chunk, render_kwargs_test)

print('Done, saving', rgbs.shape, disps.shape)

moviebase = os.path.join(basedir, expname, '{}_spiral_{:06d}_'.format(expname, i))

imageio.mimwrite(moviebase + 'rgb.mp4', to8b(rgbs), fps=30, quality=8)

imageio.mimwrite(moviebase + 'disp.mp4', to8b(disps / np.max(disps)), fps=30, quality=8)

# if args.use_viewdirs:

# render_kwargs_test['c2w_staticcam'] = render_poses[0][:3,:4]

# with torch.no_grad():

# rgbs_still, _ = render_path(render_poses, hwf, args.chunk, render_kwargs_test)

# render_kwargs_test['c2w_staticcam'] = None

# imageio.mimwrite(moviebase + 'rgb_still.mp4', to8b(rgbs_still), fps=30, quality=8)

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_test].shape)

with torch.no_grad():

render_path(torch.Tensor(poses[i_test]).to(device), hwf, K, args.chunk, render_kwargs_test, gt_imgs=images[i_test], savedir=testsavedir)

print('Saved test set')

if i%args.i_print==0:

tqdm.write(f"[TRAIN] Iter: {i} Loss: {loss.item()} PSNR: {psnr.item()}")

"""