Gaussian Reconstruction model inputs

[mparelli]

Hi @chenguolin, @paulpanwang

Thank you for the great work! I have a question regarding the inputs of the gaussian reconstruction module (specifically depth and normals). In the GObjaverseParquetDataset, the line self._load_png(sample[f"{vid:05d}_nd.png is included. However, in the Gobjaverse dataset I find the nd.exr files. How did you process the nd.exr files for this specific use case?

[chenguolin]

Hi @mparelli, yes, we preprocess nd.exr for convenience.

Here is the preprocessing code of nd.exr:

os.environ["OPENCV_IO_ENABLE_OPENEXR"] = "1"
nd_path = os.path.join(view_dir, f"{i:05}_nd.exr")

nd = cv2.imread(nd_path, cv2.IMREAD_UNCHANGED)
nd[:, :, :3] = nd[:, :, :3][..., ::-1]  # BGR -> RGB
nd[:, :, :3] = (nd[:, :, :3] * 0.5 + 0.5) * 65535  # [-1., 1.] -> [0, 65535]
nd[:, :, 3] = nd[:, :, 3] / 5. * 65535  # scale the depth by 1/5, then it must be in [0, 1]; [0., +?) -> [0, 65535]
outputs[f"{i:05}_nd.png"] = tf.io.encode_png(tf.convert_to_tensor(nd, tf.uint16)).numpy()  # save bytes for parquet

Note that GSRecon is NOT really used during inference. Only its rendering function is used for visualization.

[mparelli]

Thank you for your reply @chenguolin ! Do you also preprocess the input images and camera json files of the G-Objaverse dataset before loading them in GObjaverseParquetDataset? I want to feed multiview images to the Gaussian Reconstruction model and visualize the renderings.

[chenguolin]

I've added this information in README. Please refer to here.

[mparelli]

Hi @chenguolin ,

Thank you very much for the additions! After visualizing the renderings from the Gaussian Reconstruction module (with 4 rgb/normals from G-Objaverse as input), I have some unexpected results. Could you provide a short code snippet that shows the correct way to pass inputs to the gsrecon model and output the rendered images?

[chenguolin]

Please refer to compute_loss() of GSReconhere.

[mparelli]

Thank you for your reply @chenguolin. Unfortunately, I get similar results as above. I provide my code here.

gsrecon = GSRecon(opt)
gsrecon.requires_grad_(False)
gsrecon.eval()
print(f"Load GSRecon checkpoint from [{args.load_pretrained_gsrecon}] iteration [{args.load_pretrained_gsrecon_ckpt:06d}]\n")
gsrecon = util.load_ckpt(
    os.path.join(args.output_dir, args.load_pretrained_gsrecon, "checkpoints"),
    args.load_pretrained_gsrecon_ckpt,
    None if args.hdfs_dir is None else os.path.join(args.project_hdfs_dir, args.load_pretrained_gsrecon),
    gsrecon,
)
gsrecon = gsrecon.to(f"cuda:{args.gpu_id}")
V_in = opt.num_input_views
fxfycxcy = torch.tensor([opt.fxfy, opt.fxfy, 0.5, 0.5], dtype=torch.float32) 
return_dict = defaultdict(list)
orig_path = "/path/richdreamer/dataset/gobjaverse/saved_plants/105/538974/"
init_azi = None
sample = {}
for vid in [0,6,12,18]:
    path = orig_path + f"{vid:05d}"
    json_path = path + f"/{vid:05d}.json"
    return_dict["fxfycxcy"].append(fxfycxcy)
    sample[f"{vid:05}.png"] = tf.io.encode_png(tf.convert_to_tensor(imageio.imread(path + f"/{vid:05d}.png"), tf.uint8)).numpy()
    image = _load_png(sample[f"{vid:05d}.png"], uint16=False)
    mask = image[3:4]  # (1, 512, 512)
    image = image[:3] * mask + (1. - mask)  # (3, 512, 512), to white bg
    return_dict["image"].append(image)  # (V, 3, H, W)
    return_dict["mask"].append(mask)  # (V, 1, H, W)
    with open(json_path, "r") as f:
        json_data = f.read().encode("utf-8")  # Read as string
    c2w = _load_camera_from_json(json_data)
    # Blender world + OpenCV cam -> OpenGL world & cam; https://kit.kiui.moe/camera
    c2w[1] *= -1
    c2w[[1, 2]] = c2w[[2, 1]]
    c2w[:3, 1:3] *= -1  # invert up and forward direction
    return_dict["original_C2W"].append(torch.from_numpy(c2w).float())  # (V, 4, 4); for normal normalization only
    cond_cam_dist = np.linalg.norm(c2w[:3, 3:], 2)
    # # # Relative azimuth w.r.t. the first view
    ele, azi, dis = undo_orbit_camera(c2w)  # elevation: [-90, 90] from +y(-90) to -y(90)
    if init_azi is None:
        init_azi = azi
    azi = (azi - init_azi) % 360.  # azimuth: [0, 360] from +z(0) to +x(90)
    # To avoid numerical errors for elevation +/- 90 (GObjaverse index 25 (up) & 26 (down))
    ele_sign = ele >= 0
    ele = abs(ele) - 1e-8
    ele = ele * (1. if ele_sign else -1.)
    new_c2w = torch.from_numpy(orbit_camera(ele, azi, dis)).float()
    return_dict["C2W"].append(new_c2w)  # (V, 4, 4)
    return_dict["cam_pose"].append(torch.tensor([np.deg2rad(ele), np.deg2rad(azi), dis], dtype=torch.float32))  # (V, 3)

    if opt.load_normal or opt.load_coord:
        nd_path = f"{path}/{vid:05d}_nd.exr"
        nd = cv2.imread(nd_path, cv2.IMREAD_UNCHANGED)
        nd[:, :, :3] = nd[:, :, :3][..., ::-1]  # BGR -> RGB
        nd[:, :, :3] = (nd[:, :, :3] * 0.5 + 0.5) * 65535  # [-1., 1.] -> [0, 65535]
        nd[:, :, 3] = nd[:, :, 3] / 5. * 65535  # scale the depth by 1/5, then it must be in [0, 1]; [0., +?) -> [0, 65535]
        sample[f"{vid:05}_nd.png"] = tf.io.encode_png(tf.convert_to_tensor(nd, tf.uint16)).numpy()  # save bytes for parquet
        nd = _load_png(sample[f"{vid:05d}_nd.png"], uint16=True)
        if opt.load_normal:
            normal = nd[:3] * 2. - 1.  # (3, 512, 512) in [-1, 1]
            normal[0, ...] *= -1  # to OpenGL world convention
            return_dict["normal"].append(normal)  # (V, 3, H, W)
        if opt.load_coord or opt.load_depth:
            depth = nd[3]*5  # (512, 512); NOTE: depth is scaled by 1/5 in my data preprocessing
            return_dict["depth"].append(depth)  # (V, H, W)

for key in return_dict.keys():
    return_dict[key] = torch.stack(return_dict[key], dim=0)

if opt.load_normal:
    # Normalize normals by the first view and transform the first view to a fixed azimuth (i.e., 0)
    # Ensure `normals` and `original_C2W` are in the same camera convention
    normals = normalize_normals(return_dict["normal"].unsqueeze(0), return_dict["original_C2W"].unsqueeze(0), i=0).squeeze(0)
    normals = torch.einsum("brc,bvchw->bvrhw", return_dict["C2W"][0, :3, :3].unsqueeze(0), normals.unsqueeze(0)).squeeze(0)
    normals = normals * 0.5 + 0.5  # [0, 1]
    normals = normals * return_dict["mask"] + (1. - return_dict["mask"])  # (V, 3, 512, 512), to white bg
    return_dict["normal"] = normals
    return_dict.pop("original_C2W")  # original C2W is only used for normal normalization

# OpenGL to COLMAP camera for Gaussian renderer
return_dict["C2W"][:, :3, 1:3] *= -1
if opt.norm_camera:
    scale = opt.norm_radius / (torch.norm(return_dict["C2W"][0, :3, 3], dim=-1))
else:
    scale = 1.
return_dict["C2W"][:, :3, 3] *= scale
return_dict["cam_pose"][:, 2] *= scale
if opt.load_coord:
    # Unproject depth map to 3D world coordinate
    coords = unproject_depth(return_dict["depth"].unsqueeze(0) * scale,
        return_dict["C2W"].unsqueeze(0), return_dict["fxfycxcy"].unsqueeze(0)).squeeze(0)
    coords = coords * 0.5 + 0.5  # [0, 1]
    coords = coords * return_dict["mask"] + (1. - return_dict["mask"])  # (V, 3, 512, 512), to white bg
    return_dict["coord"] = coords
    if not opt.load_depth:
        return_dict.pop("depth")
# Resize to the input resolution
for key in ["image", "mask", "albedo", "normal", "coord", "depth", "mr", "canny"]:
    if key in return_dict.keys():
        return_dict[key] = tF.interpolate(
            return_dict[key], size=(opt.input_res, opt.input_res),
                mode="bilinear", align_corners=False, antialias=True
            )
    # Handle anti-aliased normal, coord and depth (GObjaverse renders anti-aliased normal & depth)
    if opt.load_normal:
        return_dict["normal"] = return_dict["normal"] * return_dict["mask"] + (1. - return_dict["mask"])
    if opt.load_coord:
        return_dict["coord"] = return_dict["coord"] * return_dict["mask"] + (1. - return_dict["mask"])
    if opt.load_depth:
        return_dict["depth"] = return_dict["depth"] * return_dict["mask"] + (1. - return_dict["mask"])
data = dict(return_dict)
dtype = torch.float32
device = "cuda:0"
images = data["image"].to(dtype).unsqueeze(0).to(device)  # (B, V, 3, H, W)
masks = data["mask"].to(dtype).unsqueeze(0).to(device)  # (B, V, 1, H, W)
C2W = data["C2W"].to(dtype).unsqueeze(0).to(device)  # (B, V, 4, 4)
fxfycxcy = data["fxfycxcy"].to(dtype).unsqueeze(0).to(device)  # (B, V, 4)
# Input views
V_in = opt.num_input_views
input_images = images[:, :V_in, ...]
input_C2W = C2W[:, :V_in, ...]
input_fxfycxcy = fxfycxcy[:, :V_in, ...]
if opt.input_normal:
    input_images = torch.cat([input_images, data["normal"].to(device).unsqueeze(0)[:, :V_in, ...]], dim=2)
if opt.input_coord:
    input_images = torch.cat([input_images, data["coord"].to(device).unsqueeze(0)[:, :V_in, ...]], dim=2)
if opt.input_mr:
    input_images = torch.cat([input_images, data["mr"].to(device).unsqueeze(0)[:, :V_in, :2]], dim=2)
model_outputs = gsrecon.forward_gaussians(input_images, input_C2W, input_fxfycxcy)
render_outputs = gsrecon.gs_renderer.render(model_outputs, input_C2W, input_fxfycxcy, C2W, fxfycxcy)
render_images = render_outputs["image"].to(dtype)  # (B, V, 3, H, W)
render_masks = render_outputs["alpha"]  # (B, V, 1, H, W)
render_coords = render_outputs["coord"].to(dtype)  # (B, V, 3, H, W)
render_normals = render_outputs["normal"]  # (B, V, 3, H, W)
from PIL import Image
render_images = render_outputs["image"].squeeze(0)
for i, img in enumerate(render_images):
    img = (img*255.).detach().cpu().numpy().astype(np.uint8).transpose(1, 2, 0)  # Convert (3, H, W) → (H, W, 3)
    Image.fromarray(img).save(f"view_{i:02d}.png")

[chenguolin]

Some code or model weights of GSRecon appear to be incorrect. Thank you for pointing this out, and I sincerely apologize for the issue! I’ll work on resolving it as soon as possible.

---

Update:

I have fixed the code, see the difference. The bug is from the different form of rearrange during the original training:
(rearrange(x, "b c (h p1) (w p2) -> b (h w) (p1 p2 c)") ✅)
and the bug released code:
(rearrange(x, "b c (h p1) (w p2) -> b (h w) (c p1 p2)" ❌).

Please pull the latest repo.

Thank you again for pointing this bug out!

[mparelli]

Thanks for the very quick reply! Everything seems to be correct now.