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

""" # [N_rays, N_samples, 3] -> [N_rays * N_samples,3] : 샘플린된 모든 point 들 (world 기준 실제 위치인거지)

inputs_flat = torch.reshape(inputs, [-1, inputs.shape[-1]]) # pts : ray에서 샘플링된 point들, (num ray*ray당 point 수,3)

embedded = embed_fn(inputs_flat) # [N_rays * N_samples,63]미리 틀 만들어놓은 postional encoding 거치고

if viewdirs is not None:

input_dirs = viewdirs[:, None].expand(

inputs.shape) # [1024,3]->[1024,1,3]->[1024,64,3] : 동일한 ray의 viewdir은 같으니까 차원늘려 확장만 시킨 것

input_dirs_flat = torch.reshape(input_dirs, [-1, input_dirs.shape[-1]]) # [1024*64,3]

embedded_dirs = embeddirs_fn(input_dirs_flat) # [1024*64,27] view dir positional encoding

embedded = torch.cat([embedded, embedded_dirs], -1) # [1024*64,90]

outputs_flat = batchify(fn, netchunk)(embedded) # [65536,4]본격 NeRF network forward 진입 후

outputs = torch.reshape(outputs_flat, list(inputs.shape[:-1]) + [outputs_flat.shape[-1]]) # [1024,64,4]

"""Render rays in smaller minibatches to avoid OOM(out of memory).

all_ret = {}

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

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 , test 때 넘어온 c2w

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

else: # None

# 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, K, c2w_staticcam)

viewdirs = viewdirs / torch.norm(viewdirs, dim=-1,

keepdim=True) # rays_d를 unit-magnitude viewing direction으로 만든

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() # (1024,3)

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]) # ray=o+td에서 td near,far 파트인가???

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}

render_factor=0):

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)

# TODO:ADD Depth image

depth_filename = os.path.join(savedir_depth, '{:03d}.png'.format(i))

depth8 = to8b(disps[-1] / np.max(disps[-1])) # TODO : shape check error why???

imageio.imwrite(depth_filename, depth8)

"""

rgbs = np.stack(rgbs, 0)

disps = np.stack(disps, 0)

"""Instantiate NeRF's MLP model.

# positional encoding : xyz position

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

# multires = 10, positional encodeing's L, viewing direction'L = multires_views

# positional encoding : viewing direction

input_ch_views = 0

embeddirs_fn = None

if args.use_viewdirs:

embeddirs_fn, input_ch_views = get_embedder(args.multires_views, args.i_embed)

output_ch = 5 if args.N_importance > 0 else 4 # N_importance : number of additional fine samples per ray

skips = [4]

""" input parameters

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: # 128, fine network 한다면 모델 하나 더 만들어

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, # embed position

embeddirs_fn=embeddirs_fn,

# embed view direction

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: # None일때

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, # run_network

'perturb': args.perturb, # 1.0

'N_importance': args.N_importance, # 128,fine network일때 ray에서 샘플링 수

'network_fine': model_fine, # fine NeRF network

'N_samples': args.N_samples, # number of coarse samples per ray

'network_fn': model,

'use_viewdirs': args.use_viewdirs, # view direction 하냐 True

'white_bkgd': args.white_bkgd,

'raw_noise_std': args.raw_noise_std, # 0.0 , std dev of noise added to regularize sigma_a output

}

# 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) # wi식(5), alpha 식

dists = z_vals[..., 1:] - z_vals[..., :-1] # (1024,63) : sampling point 사이 간격

dists = torch.cat([dists, torch.Tensor([1e10]).expand(dists[..., :1].shape)],

-1) # [N_rays, N_samples] [1024,64]: 위에 사이 간격에서 맨 마지막에 1e10 값 넣어

# TODO : 이해하기 rays_d(1024,3)->(1024,1,3), (1024,64)*(1024,1,3)??->(1024,64): ray point에 direction이 반영된??

dists = dists * torch.norm(rays_d[..., None, :], dim=-1) # (1024,64) 곱 연산(element-wise), norm 연산 뭐지 ,

# 아무툰 dist : ray direction 반영된 point들 맨 마지막꺼는 아래 연산할때 제외하니까 상관없어

rgb = torch.sigmoid(raw[..., :3]) # [N_rays, N_samples, 3], network에서 뒷부분에 sigmoid 통과 파트

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](1024,3)

"""

depth_map = torch.sum(weights * z_vals, -1) # (1024,64)*(1024,64) shape같으니 같은 위치 원소들끼리 곱해짐, 5번식

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

""" Disparity map. Inverse of depth map

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.

# ray_batch에서 rays_o,rays_d,viewdirs(rays_d를 유닛 크기로한 view direction),near,far 분리

N_rays = ray_batch.shape[0] # ray에서 샘플링 표본 수

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

near, far = bounds[..., 0], bounds[..., 1] # [-1,1]

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

if not lindisp: # depth에서 linaer하게 표본 추출

z_vals = near * (1. - t_vals) + far * (t_vals) # (1024,64) , 노트 참고 near---t(near,far사이 t 지점)---far

else:

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

z_vals = z_vals.expand([N_rays, N_samples]) # 각 1024가 2~6 사이 균등하게 64개 값

if perturb > 0.: # 0이 아니라면 계층화된 랜덤 포인트에서 샘플링 : sec 5.2 참고

# get intervals between samples

mids = .5 * (z_vals[..., 1:] + z_vals[..., :-1]) # (1024,63) 64개 포인트 사이 63개 구간의 가운데 점

upper = torch.cat([mids, z_vals[..., -1:]], -1) # mid와 마지막 지점 far

lower = torch.cat([z_vals[..., :1], mids], -1) # near와 mid 지점

# stratified samples in those intervals

t_rand = torch.rand(z_vals.shape) # (1024,64)

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

# [num_rays, num_samples along ray]. Integration time.

z_vals = lower + (upper - lower) * t_rand # 각 mid point 사이 내 랜덤 지점 셀렉(시작,끝 지점은 시작,끝 지점과 mid point 사이 중 랜덤 추출)

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

# ray에서 샘플링된 point 들, None = np.newaxis랑 같음

# raw = run_network(pts) , coarse network인가

raw = network_query_fn(pts, viewdirs,

network_fn) # (1024,64,4), [num_rays, num_samples along ray, 4]. Prediction from model. , run_network로 진입

rgb_map, disp_map, acc_map, weights, depth_map = raw2outputs(raw, z_vals, rays_d, raw_noise_std, white_bkgd,

pytest=pytest)

# acc_map : Sum of weights along each ray,

# fine network

if N_importance > 0:

# coarse results

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]) # random하게 뽑은 포인트들 중점

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) # (num of ray, coarse+fine)

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 = run_network(pts, fn=run_fn)

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: # coarse 결과도 따로 저장

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

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

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

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

## tum flags

# parser.add_argument("--train_sample_ratio", type=float, default=0.6,

# help='will take 6/10 images as tum test set')

# parser.add_argument("--test_sample_ratio", type=float, default=0.2,

# help='will take 6/10 images as tum test set')

# parser.add_argument("--val_sample_ratio", type=float, default=0.2,

# help='will take 6/10 images as tum test set')

# parser.add_argument("--tumhold", type=int, default=8,

# help='will take every 1/N images as LLFF test set, paper uses 8')

parser.add_argument("--tum_num_keyframe", type=int, default=120,

help='num key frame')

## arkit flags

parser.add_argument("--arkit_num_keyframe", type=int, default=120,

help='num key frame')

parser.add_argument("--min_angle_keyframe", type=float, default=15,

help='minimum angle between key frames')

parser.add_argument("--min_distance_keyframe", type=float, default=0.1,

help='minimum distance between key frames')

# 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':

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

recenter=True, bd_factor=.75,

spherify=args.spherify)

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]

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)

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:]) # 투명인 곳 흰색으로 만든다.?? 컬러 배치 확인해야해 rgba인지 뭔지

else:

images = images[..., :3]

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]

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.

# TUM dataset

elif args.dataset_type == 'tum':

images, poses, render_poses, hwf, K, i_split = load_tum_data(basedir=args.datadir,

tum_num_keyframe=args.tum_num_keyframe)

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

i_train, i_val, i_test = i_split

near = 0.5

far = 4.

if args.white_bkgd: # 물체 흰색 배경에 합성하는지

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

1. - images[..., -1:]) # 투명인 곳 흰색으로 만든다.?? 컬러 배치 확인해야해 rgba인지 뭔지

else:

images = images[..., :3]

# ARKit dataset

elif args.dataset_type == 'arkit':

images, poses, render_poses, hwf, K, i_split = load_arkit_data(basedir=args.datadir,

#arkit_num_keyframe=args.tum_num_keyframe,

min_angle=args.min_angle_keyframe,

min_distance=args.min_distance_keyframe)

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

i_train, i_val, i_test = i_split

near = 0.5

far = 6.

if args.white_bkgd: # 물체 흰색 배경에 합성하는지

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

1. - images[..., -1:]) # 투명인 곳 흰색으로 만든다.?? 컬러 배치 확인해야해 rgba인지 뭔지

else:

images = images[..., :3]

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]

if K is None:

K = np.array([

[focal, 0, 0.5 * W],

[0, focal, 0.5 * H],

[0, 0, 1]

])

if args.render_test: # render pose path 대신 test set으로 렌더링

render_poses = np.array(poses[i_test])

# Create log dir and copy the config file, logs에 폴더 생성

basedir = args.basedir # ./logs

expname = args.expname

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

# TODO : test_pose.txt save code, for depth to point cloud, after making exp directory

# test pose -> .txt save

test_pose_savedir = os.path.join(args.basedir, args.expname, 'test_poses')

np.save(test_pose_savedir, poses[i_test])

# Create nerf model

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

global_step = start

bds_dict = {

'near': near, # 2

'far': far, # 6

}

render_kwargs_train.update(bds_dict) # near,far 추가

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 = args.N_rand # 1024=2^10(L), batch size, number of random rays per gradient step , (section 6.4)

use_batching = not args.no_batching

if use_batching:

# For random ray batching

print('get rays')

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

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

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

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

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)

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

for i in trange(start, N_iters): # 1~200000, trange : progress bar

time0 = time.time()

# i = 50000 #TODO: erase

# 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:

# Random from one image

img_i = np.random.choice(i_train) # train 개수 중 랜덤 선택

target = images[img_i]

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

pose = poses[img_i, :3, :4] # (3,4), camer2world or world2camera??

if N_rand is not None: # 1024=2^10(L), batch size

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

if i < args.precrop_iters: # 500

# center cropping , [Config] Center cropping of size 200 x 200 is enabled until iter 500

dH = int(H // 2 * args.precrop_frac) # H=400 -> dh=100

dW = int(W // 2 * args.precrop_frac)

coords = torch.stack(

torch.meshgrid(

torch.linspace(H // 2 - dH, H // 2 + dH - 1, 2 * dH), # 100~299

torch.linspace(W // 2 - dW, W // 2 + dW - 1, 2 * dW)

), -1) # (200,200,2)=(H/2,W/2,2)

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 안에서 N_rand 만큼 랜덤 선택

coords = torch.reshape(coords, [-1, 2]) # (H/2 * W/2, 2) or (H,W,2)

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

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

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) , 그 해당 이미지에서 random으로 고른 pixel 값만

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

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) # img2mse = lambda x, y : torch.mean((x - y) ** 2)

trans = extras['raw'][..., -1] # raw 결과에서 투명도만

loss = img_loss # fine error

psnr = mse2psnr(img_loss) # mse2psnr = lambda x : -10. * torch.log(x) / torch.log(torch.Tensor([10.]))

# 영상 손실로 인하여 화질 손실 정보를 평가 할 때 사용, mse 적을 수록 PSNR 높음

if 'rgb0' in extras:

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

loss = loss + img_loss0 # coarse e + fine e

psnr0 = mse2psnr(img_loss0)

loss.backward()

optimizer.step()

# NOTE: IMPORTANT!

### update learning rate ###

decay_rate = 0.1

decay_steps = args.lrate_decay * 1000 # 500 * 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 #####

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

testsavedir_depth = os.path.join(basedir, expname, 'testset_depth_{:06d}'.format(i))

os.makedirs(testsavedir, exist_ok=True)

os.makedirs(testsavedir_depth, exist_ok=True)