408 lines
13 KiB
Python
408 lines
13 KiB
Python
import glob
|
|
import os
|
|
import os.path as osp
|
|
|
|
import fire
|
|
import numpy as np
|
|
import torch
|
|
import torch.nn.functional as F
|
|
from PIL import Image
|
|
from tqdm import tqdm
|
|
|
|
from seva.data_io import get_parser
|
|
from seva.eval import (
|
|
IS_TORCH_NIGHTLY,
|
|
compute_relative_inds,
|
|
create_transforms_simple,
|
|
infer_prior_inds,
|
|
infer_prior_stats,
|
|
run_one_scene,
|
|
)
|
|
from seva.geometry import (
|
|
generate_interpolated_path,
|
|
generate_spiral_path,
|
|
get_arc_horizontal_w2cs,
|
|
get_default_intrinsics,
|
|
get_lookat,
|
|
get_preset_pose_fov,
|
|
)
|
|
from seva.model import SGMWrapper
|
|
from seva.modules.autoencoder import AutoEncoder
|
|
from seva.modules.conditioner import CLIPConditioner
|
|
from seva.sampling import DDPMDiscretization, DiscreteDenoiser
|
|
from seva.utils import load_model
|
|
|
|
device = "cuda:0"
|
|
|
|
|
|
# Constants.
|
|
WORK_DIR = "work_dirs/demo"
|
|
|
|
if IS_TORCH_NIGHTLY:
|
|
COMPILE = True
|
|
os.environ["TORCHINDUCTOR_AUTOGRAD_CACHE"] = "1"
|
|
os.environ["TORCHINDUCTOR_FX_GRAPH_CACHE"] = "1"
|
|
else:
|
|
COMPILE = False
|
|
|
|
MODEL = SGMWrapper(load_model(device="cpu", verbose=True).eval()).to(device)
|
|
AE = AutoEncoder(chunk_size=1).to(device)
|
|
CONDITIONER = CLIPConditioner().to(device)
|
|
DISCRETIZATION = DDPMDiscretization()
|
|
DENOISER = DiscreteDenoiser(discretization=DISCRETIZATION, num_idx=1000, device=device)
|
|
VERSION_DICT = {
|
|
"H": 576,
|
|
"W": 576,
|
|
"T": 21,
|
|
"C": 4,
|
|
"f": 8,
|
|
"options": {},
|
|
}
|
|
|
|
if COMPILE:
|
|
MODEL = torch.compile(MODEL, dynamic=False)
|
|
CONDITIONER = torch.compile(CONDITIONER, dynamic=False)
|
|
AE = torch.compile(AE, dynamic=False)
|
|
|
|
|
|
def parse_task(
|
|
task,
|
|
scene,
|
|
num_inputs,
|
|
T,
|
|
version_dict,
|
|
):
|
|
options = version_dict["options"]
|
|
|
|
anchor_indices = None
|
|
anchor_c2ws = None
|
|
anchor_Ks = None
|
|
|
|
if task == "img2trajvid_s-prob":
|
|
if num_inputs is not None:
|
|
assert (
|
|
num_inputs == 1
|
|
), "Task `img2trajvid_s-prob` only support 1-view conditioning..."
|
|
else:
|
|
num_inputs = 1
|
|
num_targets = options.get("num_targets", T - 1)
|
|
num_anchors = infer_prior_stats(
|
|
T,
|
|
num_inputs,
|
|
num_total_frames=num_targets,
|
|
version_dict=version_dict,
|
|
)
|
|
|
|
input_indices = [0]
|
|
anchor_indices = np.linspace(1, num_targets, num_anchors).tolist()
|
|
|
|
all_imgs_path = [scene] + [None] * num_targets
|
|
|
|
c2ws, fovs = get_preset_pose_fov(
|
|
option=options.get("traj_prior", "orbit"),
|
|
num_frames=num_targets + 1,
|
|
start_w2c=torch.eye(4),
|
|
look_at=torch.Tensor([0, 0, 10]),
|
|
)
|
|
|
|
with Image.open(scene) as img:
|
|
W, H = img.size
|
|
aspect_ratio = W / H
|
|
Ks = get_default_intrinsics(fovs, aspect_ratio=aspect_ratio) # unormalized
|
|
Ks[:, :2] *= (
|
|
torch.tensor([W, H]).reshape(1, -1, 1).repeat(Ks.shape[0], 1, 1)
|
|
) # normalized
|
|
Ks = Ks.numpy()
|
|
|
|
anchor_c2ws = c2ws[[round(ind) for ind in anchor_indices]]
|
|
anchor_Ks = Ks[[round(ind) for ind in anchor_indices]]
|
|
|
|
else:
|
|
parser = get_parser(
|
|
parser_type="reconfusion",
|
|
data_dir=scene,
|
|
normalize=False,
|
|
)
|
|
all_imgs_path = parser.image_paths
|
|
c2ws = parser.camtoworlds
|
|
camera_ids = parser.camera_ids
|
|
Ks = np.concatenate([parser.Ks_dict[cam_id][None] for cam_id in camera_ids], 0)
|
|
|
|
if num_inputs is None:
|
|
assert len(parser.splits_per_num_input_frames.keys()) == 1
|
|
num_inputs = list(parser.splits_per_num_input_frames.keys())[0]
|
|
split_dict = parser.splits_per_num_input_frames[num_inputs] # type: ignore
|
|
elif isinstance(num_inputs, str):
|
|
split_dict = parser.splits_per_num_input_frames[num_inputs] # type: ignore
|
|
num_inputs = int(num_inputs.split("-")[0]) # for example 1_from32
|
|
else:
|
|
split_dict = parser.splits_per_num_input_frames[num_inputs] # type: ignore
|
|
|
|
num_targets = len(split_dict["test_ids"])
|
|
|
|
if task == "img2img":
|
|
# Note in this setting, we should refrain from using all the other camera
|
|
# info except ones from sampled_indices, and most importantly, the order.
|
|
num_anchors = infer_prior_stats(
|
|
T,
|
|
num_inputs,
|
|
num_total_frames=num_targets,
|
|
version_dict=version_dict,
|
|
)
|
|
|
|
sampled_indices = np.sort(
|
|
np.array(split_dict["train_ids"] + split_dict["test_ids"])
|
|
) # we always sort all indices first
|
|
|
|
traj_prior = options.get("traj_prior", None)
|
|
if traj_prior == "spiral":
|
|
assert parser.bounds is not None
|
|
anchor_c2ws = generate_spiral_path(
|
|
c2ws[sampled_indices] @ np.diagflat([1, -1, -1, 1]),
|
|
parser.bounds[sampled_indices],
|
|
n_frames=num_anchors + 1,
|
|
n_rots=2,
|
|
zrate=0.5,
|
|
endpoint=False,
|
|
)[1:] @ np.diagflat([1, -1, -1, 1])
|
|
elif traj_prior == "interpolated":
|
|
assert num_inputs > 1
|
|
anchor_c2ws = generate_interpolated_path(
|
|
c2ws[split_dict["train_ids"], :3],
|
|
round((num_anchors + 1) / (num_inputs - 1)),
|
|
endpoint=False,
|
|
)[1 : num_anchors + 1]
|
|
elif traj_prior == "orbit":
|
|
c2ws_th = torch.as_tensor(c2ws)
|
|
lookat = get_lookat(
|
|
c2ws_th[sampled_indices, :3, 3],
|
|
c2ws_th[sampled_indices, :3, 2],
|
|
)
|
|
anchor_c2ws = torch.linalg.inv(
|
|
get_arc_horizontal_w2cs(
|
|
torch.linalg.inv(c2ws_th[split_dict["train_ids"][0]]),
|
|
lookat,
|
|
-F.normalize(
|
|
c2ws_th[split_dict["train_ids"]][:, :3, 1].mean(0),
|
|
dim=-1,
|
|
),
|
|
num_frames=num_anchors + 1,
|
|
endpoint=False,
|
|
)
|
|
).numpy()[1:, :3]
|
|
else:
|
|
anchor_c2ws = None
|
|
# anchor_Ks is default to be the first from target_Ks
|
|
|
|
all_imgs_path = [all_imgs_path[i] for i in sampled_indices]
|
|
c2ws = c2ws[sampled_indices]
|
|
Ks = Ks[sampled_indices]
|
|
|
|
# absolute to relative indices
|
|
input_indices = compute_relative_inds(
|
|
sampled_indices,
|
|
np.array(split_dict["train_ids"]),
|
|
)
|
|
anchor_indices = np.arange(
|
|
sampled_indices.shape[0],
|
|
sampled_indices.shape[0] + num_anchors,
|
|
).tolist() # the order has no meaning here
|
|
|
|
elif task == "img2vid":
|
|
num_targets = len(all_imgs_path) - num_inputs
|
|
num_anchors = infer_prior_stats(
|
|
T,
|
|
num_inputs,
|
|
num_total_frames=num_targets,
|
|
version_dict=version_dict,
|
|
)
|
|
|
|
input_indices = split_dict["train_ids"]
|
|
anchor_indices = infer_prior_inds(
|
|
c2ws,
|
|
num_prior_frames=num_anchors,
|
|
input_frame_indices=input_indices,
|
|
options=options,
|
|
).tolist()
|
|
num_anchors = len(anchor_indices)
|
|
anchor_c2ws = c2ws[anchor_indices, :3]
|
|
anchor_Ks = Ks[anchor_indices]
|
|
|
|
elif task == "img2trajvid":
|
|
num_anchors = infer_prior_stats(
|
|
T,
|
|
num_inputs,
|
|
num_total_frames=num_targets,
|
|
version_dict=version_dict,
|
|
)
|
|
|
|
target_c2ws = c2ws[split_dict["test_ids"], :3]
|
|
target_Ks = Ks[split_dict["test_ids"]]
|
|
anchor_c2ws = target_c2ws[
|
|
np.linspace(0, num_targets - 1, num_anchors).round().astype(np.int64)
|
|
]
|
|
anchor_Ks = target_Ks[
|
|
np.linspace(0, num_targets - 1, num_anchors).round().astype(np.int64)
|
|
]
|
|
|
|
sampled_indices = split_dict["train_ids"] + split_dict["test_ids"]
|
|
all_imgs_path = [all_imgs_path[i] for i in sampled_indices]
|
|
c2ws = c2ws[sampled_indices]
|
|
Ks = Ks[sampled_indices]
|
|
|
|
input_indices = np.arange(num_inputs).tolist()
|
|
anchor_indices = np.linspace(
|
|
num_inputs, num_inputs + num_targets - 1, num_anchors
|
|
).tolist()
|
|
|
|
else:
|
|
raise ValueError(f"Unknown task: {task}")
|
|
|
|
return (
|
|
all_imgs_path,
|
|
num_inputs,
|
|
num_targets,
|
|
input_indices,
|
|
anchor_indices,
|
|
torch.tensor(c2ws[:, :3]).float(),
|
|
torch.tensor(Ks).float(),
|
|
(torch.tensor(anchor_c2ws[:, :3]).float() if anchor_c2ws is not None else None),
|
|
(torch.tensor(anchor_Ks).float() if anchor_Ks is not None else None),
|
|
)
|
|
|
|
|
|
def main(
|
|
data_path,
|
|
data_items=None,
|
|
task="img2img",
|
|
save_subdir="",
|
|
H=None,
|
|
W=None,
|
|
T=None,
|
|
use_traj_prior=False,
|
|
**overwrite_options,
|
|
):
|
|
if H is not None:
|
|
VERSION_DICT["H"] = H
|
|
if W is not None:
|
|
VERSION_DICT["W"] = W
|
|
if T is not None:
|
|
VERSION_DICT["T"] = [int(t) for t in T.split(",")] if isinstance(T, str) else T
|
|
|
|
options = VERSION_DICT["options"]
|
|
options["chunk_strategy"] = "nearest-gt"
|
|
options["video_save_fps"] = 30.0
|
|
options["beta_linear_start"] = 5e-6
|
|
options["log_snr_shift"] = 2.4
|
|
options["guider_types"] = 1
|
|
options["cfg"] = 2.0
|
|
options["camera_scale"] = 2.0
|
|
options["num_steps"] = 50
|
|
options["cfg_min"] = 1.2
|
|
options["encoding_t"] = 1
|
|
options["decoding_t"] = 1
|
|
options["num_inputs"] = None
|
|
options["seed"] = 23
|
|
options.update(overwrite_options)
|
|
|
|
num_inputs = options["num_inputs"]
|
|
seed = options["seed"]
|
|
|
|
if data_items is not None:
|
|
if not isinstance(data_items, (list, tuple)):
|
|
data_items = data_items.split(",")
|
|
scenes = [os.path.join(data_path, item) for item in data_items]
|
|
else:
|
|
scenes = glob.glob(osp.join(data_path, "*"))
|
|
|
|
for scene in tqdm(scenes):
|
|
save_path_scene = os.path.join(
|
|
WORK_DIR, task, save_subdir, os.path.splitext(os.path.basename(scene))[0]
|
|
)
|
|
if options.get("skip_saved", False) and os.path.exists(
|
|
os.path.join(save_path_scene, "transforms.json")
|
|
):
|
|
print(f"Skipping {scene} as it is already sampled.")
|
|
continue
|
|
|
|
# parse_task -> infer_prior_stats modifies VERSION_DICT["T"] in-place.
|
|
(
|
|
all_imgs_path,
|
|
num_inputs,
|
|
num_targets,
|
|
input_indices,
|
|
anchor_indices,
|
|
c2ws,
|
|
Ks,
|
|
anchor_c2ws,
|
|
anchor_Ks,
|
|
) = parse_task(
|
|
task,
|
|
scene,
|
|
num_inputs,
|
|
VERSION_DICT["T"],
|
|
VERSION_DICT,
|
|
)
|
|
assert num_inputs is not None
|
|
# Create image conditioning.
|
|
image_cond = {
|
|
"img": all_imgs_path,
|
|
"input_indices": input_indices,
|
|
"prior_indices": anchor_indices,
|
|
}
|
|
# Create camera conditioning.
|
|
camera_cond = {
|
|
"c2w": c2ws.clone(),
|
|
"K": Ks.clone(),
|
|
"input_indices": list(range(num_inputs + num_targets)),
|
|
}
|
|
# run_one_scene -> transform_img_and_K modifies VERSION_DICT["H"] and VERSION_DICT["W"] in-place.
|
|
video_path_generator = run_one_scene(
|
|
task,
|
|
VERSION_DICT, # H, W maybe updated in run_one_scene
|
|
model=MODEL,
|
|
ae=AE,
|
|
conditioner=CONDITIONER,
|
|
denoiser=DENOISER,
|
|
image_cond=image_cond,
|
|
camera_cond=camera_cond,
|
|
save_path=save_path_scene,
|
|
use_traj_prior=use_traj_prior,
|
|
traj_prior_Ks=anchor_Ks,
|
|
traj_prior_c2ws=anchor_c2ws,
|
|
seed=seed, # to ensure sampled video can be reproduced in regardless of start and i
|
|
)
|
|
for _ in video_path_generator:
|
|
pass
|
|
|
|
# Convert from OpenCV to OpenGL camera format.
|
|
c2ws = c2ws @ torch.tensor(np.diag([1, -1, -1, 1])).float()
|
|
img_paths = sorted(glob.glob(osp.join(save_path_scene, "samples-rgb", "*.png")))
|
|
if len(img_paths) != len(c2ws):
|
|
input_img_paths = sorted(
|
|
glob.glob(osp.join(save_path_scene, "input", "*.png"))
|
|
)
|
|
assert len(img_paths) == num_targets
|
|
assert len(input_img_paths) == num_inputs
|
|
assert c2ws.shape[0] == num_inputs + num_targets
|
|
target_indices = [i for i in range(c2ws.shape[0]) if i not in input_indices]
|
|
img_paths = [
|
|
input_img_paths[input_indices.index(i)]
|
|
if i in input_indices
|
|
else img_paths[target_indices.index(i)]
|
|
for i in range(c2ws.shape[0])
|
|
]
|
|
create_transforms_simple(
|
|
save_path=save_path_scene,
|
|
img_paths=img_paths,
|
|
img_whs=np.array([VERSION_DICT["W"], VERSION_DICT["H"]])[None].repeat(
|
|
num_inputs + num_targets, 0
|
|
),
|
|
c2ws=c2ws,
|
|
Ks=Ks,
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
fire.Fire(main)
|