def imread(fs):
if fs.endswith("png"):
with file_utils.open_file(fs) as f:
return imageio.imread(f, ignoregamma=True)
else:
with file_utils.open_file(fs) as f:
return imageio.imread(f)
def get_train_scene_list(config):
"""Function to get the list of scenes.
Args:
config: experiment config.
Returns:
scene_path_list: list of scenes.
"""
if config.dataset.name == "ff_epipolar":
corrupted_and_test_list = [
"howardzhou_010_internal_drawing_vase", "howardzhou_059_narcissus",
"howardzhou_087_yellow_chain_links",
"howardzhou_089_whilte_bmw_x3_front", "howardzhou_085_sweet_onions",
"qq18", "qq33", "data2_fernvlsb", "data2_hugetrike", "data2_trexsanta",
"data3_orchid", "data5_leafscene", "data5_lotr", "data5_redflower"
]
scene_path_list = file_utils.listdir(config.dataset.ff_base_dir)
scene_path_list = list(set(scene_path_list) - set(corrupted_and_test_list))
elif config.dataset.name == "dtu":
with file_utils.open_file(
os.path.join(config.dataset.dtu_base_dir, "configs", "lists",
"dtu_train_all.txt")) as f:
scene_path_list = [
line.rstrip().decode("utf-8") for line in f.readlines()
]
elif config.dataset.name == "blender_rot":
scene_path_list = ["lego"]
return scene_path_list
def save_img(img, pth):
"""Save an image to disk.
Args:
img: jnp.ndarry, [height, width, channels], img will be clipped to [0, 1]
before saved to pth.
pth: string, path to save the image to.
"""
with file_utils.open_file(pth, "wb") as imgout:
Image.fromarray(
np.array((np.clip(img, 0., 1.) * 255.).astype(jnp.uint8))).save(
imgout, "PNG")
def _load_renderings(self, args):
"""Load images and camera information."""
self.cam_transform = np.array([[1, 0, 0, 0], [0, -1, 0, 0],
[0, 0, -1, 0], [0, 0, 0, 1]])
#-------------------------------------------
# Load images.
#-------------------------------------------
basedir = path.join(args.dataset.eval_ff_dir, self.scene)
img0 = [
os.path.join(basedir, "images", f) for f in sorted(
file_utils.listdir(os.path.join(basedir, "images")))
if f.endswith("JPG") or f.endswith("jpg") or f.endswith("png")
][0]
with file_utils.open_file(img0) as f:
sh = imageio.imread(f).shape
if sh[0] / sh[
1] != args.dataset.eval_ff_image_height / args.dataset.eval_ff_image_width:
raise ValueError("not expected height width ratio")
factor = sh[0] / args.dataset.eval_ff_image_height
sfx = "_4"
imgdir = os.path.join(basedir, "images" + sfx)
if not file_utils.file_exists(imgdir):
imgdir = os.path.join(basedir, "images")
if not file_utils.file_exists(imgdir):
raise ValueError("{} does not exist".format(imgdir))
images = self._load_images(imgdir, args.dataset.eval_ff_image_width,
args.dataset.eval_ff_image_height)
#-------------------------------------------
# Load poses and bds.
#-------------------------------------------
with file_utils.open_file(path.join(basedir, "poses_bounds.npy"),
"rb") as fp:
poses_arr = np.load(fp)
# Get the intrinsic matrix
with file_utils.open_file(path.join(basedir, "hwf_cxcy.npy"),
"rb") as fp:
self.intrinsic_arr = np.load(fp)
# Update the intrinsic matix to accounto for resizing
self.intrinsic_arr = self.intrinsic_arr * 1. / factor
# poses_arr contains an array consisting of a 3x4 pose matrices and
# 2 depth bounds for each image. The pose matrix contain [R t] as the
# left 3x4 matrix
# pose_arr has shape (...,14) {3x4 + 2}
poses = poses_arr[:, :-2].reshape([-1, 3, 4]).transpose([1, 2, 0])
bds = poses_arr[:, -2:].transpose([1, 0])
# Convert R matrix from the form [down right back] to [right up back]
poses = np.concatenate(
[poses[:, 1:2, :], -poses[:, 0:1, :], poses[:, 2:, :]], 1)
# Transpose such that the first dimension is number of images
images = np.moveaxis(images, -1, 0)
poses = np.moveaxis(poses, -1, 0).astype(np.float32)
bds = np.moveaxis(bds, -1, 0).astype(np.float32)
if args.dataset.normalize:
scale = 1. / bds.max()
else:
scale = 1. / (bds.min() * .75)
poses[:, :3, 3] *= scale
bds *= scale
poses_copy = poses.copy()
poses_copy = pose_utils.recenter_poses(poses, None)
poses = pose_utils.recenter_poses(poses, self.cam_transform)
# Get the min and max depth of the scene
self.min_depth = np.array([bds.min()])
self.max_depth = np.array([bds.max()])
# Use this to set the near and far plane
args.model.near = self.min_depth.item()
args.model.far = self.max_depth.item()
if self.split == "test":
self.render_poses = pose_utils.generate_spiral_poses(
poses_copy, bds, self.cam_transform)
# Select the split.
if args.eval.mvsn_style:
with file_utils.open_file(
os.path.join(os.path.dirname(basedir), "pairs.npz")) as f:
img_ids = np.load(f)
i_test = img_ids["{}_{}".format(os.path.basename(basedir),
"test")]
i_train = img_ids["{}_{}".format(os.path.basename(basedir),
"train")]
else:
i_test = np.arange(images.shape[0])[::args.dataset.llffhold]
i_train = np.array([
i for i in np.arange(int(images.shape[0])) if i not in i_test
])
if self.split == "train":
indices = i_train
else:
indices = i_test
images = images[indices]
poses = poses[indices]
self.images = images
self.camtoworlds = poses[:, :3, :4]
# intrinsic arr has H, W, fx, fy, cx, cy
self.focal = self.intrinsic_arr[2][0]
self.h, self.w = images.shape[1:3]
self.resolution = self.h * self.w
if args.dataset.render_path and self.split == "test":
self.n_examples = self.render_poses.shape[0]
else:
self.n_examples = images.shape[0]
_, _, fx, fy, cx, cy = self.intrinsic_arr[:, 0]
self.intrinsic_matrix = np.array([
[fx, 0, cx, 0],
[0, fy, cy, 0],
[0, 0, 1, 0],
]).astype(np.float32)
def evaluate(config, workdir):
"""Evalution function."""
# Hide the GPUs and TPUs from TF so it does not reserve memory on them for
# LPIPS computation or dataset loading.
tf.config.experimental.set_visible_devices([], "GPU")
tf.config.experimental.set_visible_devices([], "TPU")
rng = jax.random.PRNGKey(config.seed)
if config.eval.return_coarse:
sfx = "coarse"
else:
sfx = ""
#----------------------------------------------------------------------------
# Build input pipeline.
rng, data_rng = jax.random.split(rng)
data_rng = jax.random.fold_in(data_rng, jax.process_index())
train_ds, test_ds_dict = datasets.create_eval_dataset(config)
example_batch = train_ds.peek()
rng, key = jax.random.split(rng)
#----------------------------------------------------------------------------
# Initialize model.
learning_rate_fn = train_utils.create_learning_rate_fn(config)
model, state = models.create_train_state(config,
key,
learning_rate_fn=learning_rate_fn,
example_batch=example_batch)
#----------------------------------------------------------------------------
# Get the rendering function. Renderig is forced ot be deterministic even if
# trainin is randomized
render_pfn = render_utils.get_render_function(model,
config,
randomized=False)
last_step = 0
out_dir_dict = {}
for key in test_ds_dict:
out_dir_dict[key] = os.path.join(
workdir, "path_renders" +
sfx if config.dataset.render_path else "test_preds" + sfx, key)
if not config.eval.eval_once:
summary_writer_dict = {}
for key in test_ds_dict:
# Prepare Metric Writers
summary_writer_dict[key] = tensorboard.SummaryWriter(
os.path.join(workdir, "eval" + sfx, key))
summary_writer_dict["all"] = tensorboard.SummaryWriter(
os.path.join(workdir, "eval" + sfx))
while True:
state = checkpoints.restore_checkpoint(
workdir,
state,
step=None if config.eval.checkpoint_step == -1 else
config.eval.checkpoint_step)
step = int(state.step)
if step <= last_step:
continue
total_psnr = 0
for scene_name, test_ds in test_ds_dict.items():
out_dir = out_dir_dict[scene_name]
if config.eval.save_output and (not file_utils.isdir(out_dir)):
file_utils.makedirs(out_dir)
psnr_values = []
ssim_values = []
if not config.eval.eval_once:
showcase_index = np.random.randint(0, test_ds.size)
for idx in range(test_ds.size):
logging.info("Evaluating scene %s [%d / %d].", scene_name, idx,
test_ds.size)
batch = next(test_ds)
test_pixels = batch.target_view.rgb
test_mask = batch.target_view.mask
if test_pixels is not None:
test_pixels = model_utils.uint2float(test_pixels)
#-----------------------------------------------------------
# Render Image
variables = {"params": state.params}
pred_color, pred_disp, pred_acc = render_utils.render_image(
functools.partial(render_pfn, variables),
batch,
rng,
render_utils.normalize_disp(config.dataset.name),
chunk=config.eval.chunk,
return_coarse=config.eval.return_coarse,
)
if jax.process_index() != 0:
continue
#-----------------------------------------------------------
# Get showcase example for logging
if not config.eval.eval_once and idx == showcase_index:
showcase_color = pred_color
showcase_disp = pred_disp
showcase_acc = pred_acc
if not config.dataset.render_path:
showcase_gt = test_pixels
#-----------------------------------------------------------
# If get pixels available, evaluate
if not config.dataset.render_path:
if config.eval.mvsn_style:
h_crop, w_crop = np.array(pred_color.shape[:2]) // 10
pred_color = pred_color[h_crop:-h_crop, w_crop:-w_crop]
test_pixels = test_pixels[h_crop:-h_crop,
w_crop:-w_crop]
if test_mask is not None:
psnr = model_utils.compute_psnr(
((pred_color[test_mask] -
test_pixels[test_mask])**2).mean())
else:
psnr = model_utils.compute_psnr(
((pred_color - test_pixels)**2).mean())
ssim = skmetrics.structural_similarity(
pred_color.astype(np.float32),
test_pixels.astype(np.float32),
win_size=11,
multichannel=True,
gaussian_weights=True)
logging.info(f"PSNR = {psnr:.4f}, SSIM = {ssim:.4f}") # pylint: disable=logging-fstring-interpolation
psnr_values.append(float(psnr))
ssim_values.append(float(ssim))
#-----------------------------------------------------------
# Save generated image
if config.eval.save_output:
model_utils.save_img(
pred_color,
os.path.join(out_dir, "{:03d}.png".format(idx)))
if pred_disp is not None:
model_utils.save_img(
pred_disp[Ellipsis, 0],
os.path.join(out_dir,
"disp_{:03d}.png".format(idx)))
#-----------------------------------------------------------
if (not config.eval.eval_once) and (jax.process_index() == 0):
summary_writer_dict[scene_name].image(
"eval/{}_pred_color".format(scene_name), showcase_color,
step)
if showcase_disp is not None:
summary_writer_dict[scene_name].image(
"eval/{}_pred_disp".format(scene_name), showcase_disp,
step)
if showcase_acc is not None:
summary_writer_dict[scene_name].image(
"eval/{}_pred_acc".format(scene_name), showcase_acc,
step)
if not config.dataset.render_path:
summary_writer_dict[scene_name].scalar(
"eval_metric/psnr", np.mean(np.array(psnr_values)),
step)
summary_writer_dict[scene_name].scalar(
"eval_metric/ssim", np.mean(np.array(ssim_values)),
step)
summary_writer_dict[scene_name].image(
"eval/{}_target".format(scene_name), showcase_gt, step)
#-----------------------------------------------------------
# Save the metric to file
if config.eval.save_output and (not config.dataset.render_path
) and (jax.process_index() == 0):
with file_utils.open_file(
os.path.join(out_dir, f"psnrs_{step}.txt"), "w") as f:
f.write(" ".join([str(v) for v in psnr_values]))
with file_utils.open_file(
os.path.join(out_dir, f"ssims_{step}.txt"), "w") as f:
f.write(" ".join([str(v) for v in ssim_values]))
with file_utils.open_file(os.path.join(out_dir, "psnr.txt"),
"w") as f:
f.write("{}".format(np.mean(np.array(psnr_values))))
with file_utils.open_file(os.path.join(out_dir, "ssim.txt"),
"w") as f:
f.write("{}".format(np.mean(np.array(ssim_values))))
total_psnr += np.mean(np.array(psnr_values))
if not config.eval.eval_once:
summary_writer_dict["all"].scalar(
"eval_metric/avg_psnr_{}".format(config.dataset.eval_dataset),
total_psnr / len(test_ds_dict.keys()), step)
if config.eval.eval_once:
break
if int(step) >= config.train.max_steps:
break
last_step = step
logging.info("Finishing evaluation at step %d", last_step)
def train_and_evaluate(config, workdir):
"""Runs a training and evaluation loop.
Args:
config: Configuration to use.
workdir: Working directory for checkpoints and TF summaries. If this
contains checkpoint training will be resumed from the latest checkpoint.
"""
if config.dataset.batch_size % jax.device_count() != 0:
raise ValueError(
"Batch size must be divisible by the number of devices.")
tf.io.gfile.makedirs(workdir)
# Deterministic training.
rng = jax.random.PRNGKey(config.seed)
# Shift the numpy random seed by process_index() to shuffle data loaded
# by different hosts
np.random.seed(20201473 + jax.process_index())
#----------------------------------------------------------------------------
# Build input pipeline.
rng, data_rng = jax.random.split(rng)
data_rng = jax.random.fold_in(data_rng, jax.process_index())
scene_path_list = train_utils.get_train_scene_list(config)
train_ds = datasets.create_train_dataset(config, scene_path_list[0])
_, eval_ds_dict = datasets.create_eval_dataset(config)
_, eval_ds = eval_ds_dict.popitem()
example_batch = train_ds.peek()
#----------------------------------------------------------------------------
# Learning rate schedule.
num_train_steps = config.train.max_steps
if num_train_steps == -1:
num_train_steps = train_ds.size()
steps_per_epoch = num_train_steps // config.train.num_epochs
logging.info("num_train_steps=%d, steps_per_epoch=%d", num_train_steps,
steps_per_epoch)
learning_rate_fn = train_utils.create_learning_rate_fn(config)
#----------------------------------------------------------------------------
# Initialize model.
rng, model_rng = jax.random.split(rng)
model, state = models.create_train_state(
config,
model_rng,
learning_rate_fn=learning_rate_fn,
example_batch=example_batch,
)
#----------------------------------------------------------------------------
# Set up checkpointing of the model and the input pipeline.
# check if the job was stopped and relaunced
latest_ckpt = checkpoints.latest_checkpoint(workdir)
if latest_ckpt is None:
# No previous checkpoint. Then check for pretrained weights.
if config.train.pretrain_dir:
state = checkpoints.restore_checkpoint(config.train.pretrain_dir,
state)
else:
state = checkpoints.restore_checkpoint(workdir, state)
initial_step = int(state.step) + 1
step_per_scene = config.train.switch_scene_iter
if config.dev_run:
jnp.set_printoptions(precision=2)
np.set_printoptions(precision=2)
step_per_scene = 3
#----------------------------------------------------------------------------
# Distribute training.
state = flax_utils.replicate(state)
p_train_step = jax.pmap(
functools.partial(
train_step,
model=model,
learning_rate_fn=learning_rate_fn,
weight_decay=config.train.weight_decay,
config=config,
),
axis_name="batch",
)
# Get distributed rendering function
render_pfn = render_utils.get_render_function(
model=model,
config=config,
randomized=False, # No randomization for evaluation.
)
#----------------------------------------------------------------------------
# Prepare Metric Writers
writer = metric_writers.create_default_writer(
workdir, just_logging=jax.process_index() > 0)
if initial_step == 1:
writer.write_hparams(dict(config))
logging.info("Starting training loop at step %d.", initial_step)
hooks = []
report_progress = periodic_actions.ReportProgress(
num_train_steps=num_train_steps, writer=writer)
if jax.process_index() == 0:
hooks += [
report_progress,
]
train_metrics = None
# Prefetch_buffer_size = 6 x batch_size
ptrain_ds = flax.jax_utils.prefetch_to_device(train_ds, 6)
n_local_devices = jax.local_device_count()
rng = rng + jax.process_index() # Make random seed separate across hosts.
keys = jax.random.split(rng, n_local_devices) # For pmapping RNG keys.
with metric_writers.ensure_flushes(writer):
for step in range(initial_step, num_train_steps + 1):
# `step` is a Python integer. `state.step` is JAX integer on the GPU/TPU
# devices.
if step % step_per_scene == 0:
scene_idx = np.random.randint(len(scene_path_list))
logging.info("Loading scene {}".format(
scene_path_list[scene_idx])) # pylint: disable=logging-format-interpolation
curr_scene = scene_path_list[scene_idx]
if config.dataset.name == "dtu":
# lighting can take values between 0 and 6 (both included)
config.dataset.dtu_light_idx = np.random.randint(low=0,
high=7)
train_ds = datasets.create_train_dataset(config, curr_scene)
ptrain_ds = flax.jax_utils.prefetch_to_device(train_ds, 6)
is_last_step = step == num_train_steps
with jax.profiler.StepTraceAnnotation("train", step_num=step):
batch = next(ptrain_ds)
state, metrics_update, keys = p_train_step(rng=keys,
state=state,
batch=batch)
metric_update = flax_utils.unreplicate(metrics_update)
train_metrics = (metric_update if train_metrics is None else
train_metrics.merge(metric_update))
# Quick indication that training is happening.
logging.log_first_n(logging.INFO, "Finished training step %d.", 5,
step)
for h in hooks:
h(step)
if step % config.train.log_loss_every_steps == 0 or is_last_step:
writer.write_scalars(step, train_metrics.compute())
train_metrics = None
if step % config.train.render_every_steps == 0 or is_last_step:
test_batch = next(eval_ds)
test_pixels = model_utils.uint2float(
test_batch.target_view.rgb) # extract for evaluation
with report_progress.timed("eval"):
pred_color, pred_disp, pred_acc = eval_step(
state, keys[0], test_batch, render_pfn, config)
#------------------------------------------------------------------
# Log metrics and images for host 0
#------------------------------------------------------------------
if jax.process_index() == 0:
psnr = model_utils.compute_psnr(
((pred_color - test_pixels)**2).mean())
ssim = 0.
writer.write_scalars(
step, {
"train_eval/test_psnr": psnr,
"train_eval/test_ssim": ssim,
})
writer.write_images(
step, {
"test_pred_color": pred_color[None, :],
"test_target": test_pixels[None, :]
})
if pred_disp is not None:
writer.write_images(
step, {"test_pred_disp": pred_disp[None, :]})
if pred_acc is not None:
writer.write_images(
step, {"test_pred_acc": pred_acc[None, :]})
#------------------------------------------------------------------
if (jax.process_index()
== 0) and (step % config.train.checkpoint_every_steps == 0
or is_last_step):
# Write final metrics to file
with file_utils.open_file(
os.path.join(workdir, "train_logs.json"), "w") as f:
log_dict = metric_update.compute()
for k, v in log_dict.items():
log_dict[k] = v.item()
f.write(json.dumps(log_dict))
with report_progress.timed("checkpoint"):
state_to_save = jax.device_get(
jax.tree_map(lambda x: x[0], state))
checkpoints.save_checkpoint(workdir,
state_to_save,
step,
keep=100)
logging.info("Finishing training at step %d", num_train_steps)