File size: 7,229 Bytes
73f9b1b |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 |
"""
Save a couple images to grids with cond, render cond, novel render, novel gt
Also save images to a render video
"""
import glob
import os
from PIL import Image
import numpy as np
import torch
from matplotlib import pyplot as plt
from utils.sh_utils import eval_sh
from einops import rearrange
def gridify():
out_folder = "grids_objaverse"
os.makedirs(out_folder, exist_ok=True)
folder_paths = glob.glob("/scratch/shared/beegfs/stan/scaling_splatter_image/objaverse/*")
# pixelnerf_root = "/scratch/shared/beegfs/stan/splatter_image/pixelnerf/teddybears"
folder_paths_test = sorted([fpath for fpath in folder_paths if "gt" not in fpath], key= lambda x: int(os.path.basename(x).split("_")[0]))
"""folder_paths_test = [folder_paths_test[i] for i in [5, 7, 12, 15,
18, 19, 30, 33,
37, 42, 43, 44,
48, 51, 64, 66,
70, 74, 78, 85,
89, 91, 92]]"""
# Initialize variables for grid dimensions
num_examples_row = 6
rows = num_examples_row
num_per_ex = 2
cols = num_examples_row * num_per_ex # 7 * 2
im_res = 128
for im_idx in range(100):
print("Doing frame {}".format(im_idx))
# for im_name in ["xyz", "colours", "opacity", "scaling"]:
grid = np.zeros((rows*im_res, cols*im_res, 3), dtype=np.uint8)
# Iterate through the folders in the out_folder
for f_idx, folder_path_test in enumerate(folder_paths_test[:num_examples_row*num_examples_row]):
# if im_name == "xyz":
# print(folder_path_test)
row_idx = f_idx // num_examples_row
col_idx = f_idx % num_examples_row
im_path = os.path.join(folder_path_test, "{:05d}.png".format(im_idx))
im_path_gt = os.path.join(folder_path_test + "_gt", "{:05d}.png".format(im_idx))
"""im_path_pixelnerf = os.path.join(pixelnerf_root, os.path.basename(folder_path_test),
"{:06d}.png".format(im_idx))"""
# im_path = os.path.join(folder_path_test, "{}.png".format(im_name))
try:
im = np.array(Image.open(im_path))
im_gt = np.array(Image.open(im_path_gt))
#im_pn = np.array(Image.open(im_path_pixelnerf))
grid[row_idx*im_res: (row_idx+1)*im_res,
col_idx * num_per_ex *im_res: (col_idx * num_per_ex+1)*im_res, : ] = im[:, :, :3]
grid[row_idx*im_res: (row_idx+1)*im_res,
(col_idx * num_per_ex + 1) *im_res: (col_idx* num_per_ex +2)*im_res, : ] = im_gt[:, :, :3]
"""grid[row_idx*im_res: (row_idx+1)*im_res,
(col_idx * num_per_ex + 2) *im_res: (col_idx* num_per_ex +3)*im_res, : ] = im_pn[:, :, :3]"""
except FileNotFoundError:
a = 0
im_out = Image.fromarray(grid)
im_out.save(os.path.join(out_folder, "{:05d}.png".format(im_idx)))
# im_out.save(os.path.join(out_folder, "{}.png".format(im_name)))
def comparisons():
out_root = "hydrants_comparisons"
os.makedirs(out_root, exist_ok=True)
folder_paths = glob.glob("/users/stan/pixel-nerf/full_eval_hydrant/*")
folder_paths_test = sorted(folder_paths)
folder_paths_ours_root = "/scratch/shared/beegfs/stan/out_hydrants_with_lpips_ours"
# Initialize variables for grid dimensions
rows = 3
cols = 1
im_res = 128
for f_idx, folder_path_test in enumerate(folder_paths_test):
example_id = "_".join(os.path.basename(folder_path_test).split("_")[1:])
out_folder = os.path.join(out_root, example_id)
os.makedirs(out_folder, exist_ok=True)
num_images = len([p for p in glob.glob(os.path.join(folder_path_test, "*.png")) if "gt" not in p])
grid = np.zeros((rows*im_res, cols*im_res, 3), dtype=np.uint8)
for im_idx in range(num_images):
im_path_pixelnerf = os.path.join(folder_path_test, "{:06d}.png".format(im_idx+1))
im_path_ours = os.path.join(folder_paths_ours_root, example_id, "{:05d}.png".format(im_idx))
im_path_gt = os.path.join(folder_paths_ours_root, example_id + "_gt", "{:05d}.png".format(im_idx))
# im_path = os.path.join(folder_path_test, "{}.png".format(im_name))
im_pn = np.array(Image.open(im_path_pixelnerf))
im_ours = np.array(Image.open(im_path_ours))
im_gt = np.array(Image.open(im_path_gt))
grid[:im_res, :, :] = im_pn
grid[im_res:2*im_res, :, :] = im_ours
grid[2*im_res:3*im_res, :, :] = im_gt
im_out = Image.fromarray(grid)
im_out.save(os.path.join(out_folder, "{:05d}.png".format(im_idx)))
def vis_image_preds(image_preds: dict, folder_out: str):
"""
Visualises network's image predictions.
Args:
image_preds: a dictionary of xyz, opacity, scaling, rotation, features_dc and features_rest
"""
image_preds_reshaped = {}
ray_dirs = (image_preds["xyz"].detach().cpu() / torch.norm(image_preds["xyz"].detach().cpu(), dim=-1, keepdim=True)).reshape(128, 128, 3)
for k, v in image_preds.items():
image_preds_reshaped[k] = v
if k == "xyz":
image_preds_reshaped[k] = (image_preds_reshaped[k] - torch.min(image_preds_reshaped[k], dim=0, keepdim=True)[0]) / (
torch.max(image_preds_reshaped[k], dim=0, keepdim=True)[0] - torch.min(image_preds_reshaped[k], dim=0, keepdim=True)[0]
)
if k == "scaling":
image_preds_reshaped["scaling"] = (image_preds_reshaped["scaling"] - torch.min(image_preds_reshaped["scaling"], dim=0, keepdim=True)[0]) / (
torch.max(image_preds_reshaped["scaling"], dim=0, keepdim=True)[0] - torch.min(image_preds_reshaped["scaling"], dim=0, keepdim=True)[0]
)
if k != "features_rest":
image_preds_reshaped[k] = image_preds_reshaped[k].reshape(128, 128, -1).detach().cpu()
else:
image_preds_reshaped[k] = image_preds_reshaped[k].reshape(128, 128, 3, 3).detach().cpu().permute(0, 1, 3, 2)
if k == "opacity":
image_preds_reshaped[k] = image_preds_reshaped[k].expand(128, 128, 3)
colours = torch.cat([image_preds_reshaped["features_dc"].unsqueeze(-1), image_preds_reshaped["features_rest"]], dim=-1)
colours = eval_sh(1, colours, ray_dirs)
plt.imsave(os.path.join(folder_out, "colours.png"),
colours.numpy())
plt.imsave(os.path.join(folder_out, "opacity.png"),
image_preds_reshaped["opacity"].numpy())
plt.imsave(os.path.join(folder_out, "xyz.png"),
(image_preds_reshaped["xyz"] * image_preds_reshaped["opacity"]+ 1 - image_preds_reshaped["opacity"]).numpy())
plt.imsave(os.path.join(folder_out, "scaling.png"),
(image_preds_reshaped["scaling"] * image_preds_reshaped["opacity"] + 1 - image_preds_reshaped["opacity"]).numpy())
if __name__ == "__main__":
gridify() |