mirror of
https://github.com/nerfstudio-project/nerfstudio.git
synced 2026-11-11 18:44:21 +08:00
166 lines
7.0 KiB
Python
166 lines
7.0 KiB
Python
# Copyright 2022 The Plenoptix Team. All rights reserved.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
|
|
"""
|
|
Code to handle loading friends datasets.
|
|
"""
|
|
|
|
import os
|
|
|
|
import numpy as np
|
|
import torch
|
|
|
|
from nerfactory.data.colmap_utils import (
|
|
read_cameras_binary,
|
|
read_images_binary,
|
|
read_pointsTD_binary,
|
|
)
|
|
from nerfactory.data.structs import DatasetInputs, PointCloud, SceneBounds, Semantics
|
|
from nerfactory.utils import profiler
|
|
from nerfactory.utils.io import load_from_json
|
|
|
|
|
|
def get_aabb_and_transform(basedir):
|
|
"""Returns the aabb and pointcloud transform from the threejs.json file."""
|
|
filename = os.path.join(basedir, "threejs.json")
|
|
assert os.path.exists(filename)
|
|
data = load_from_json(filename)
|
|
|
|
# point cloud transformation
|
|
transposed_point_cloud_transform = np.array(data["object"]["children"][0]["matrix"]).reshape(4, 4).T
|
|
assert transposed_point_cloud_transform[3, 3] == 1.0
|
|
|
|
# bbox transformation
|
|
bbox_transform = np.array(data["object"]["children"][1]["matrix"]).reshape(4, 4).T
|
|
w, h, d = data["geometries"][1]["width"], data["geometries"][1]["height"], data["geometries"][1]["depth"]
|
|
temp = np.array([w, h, d]) / 2.0
|
|
bbox = np.array([-temp, temp])
|
|
bbox = np.concatenate([bbox, np.ones_like(bbox[:, 0:1])], axis=1)
|
|
bbox = (bbox_transform @ bbox.T).T[:, 0:3]
|
|
|
|
aabb = bbox # rename to aabb because it's an axis-aligned bounding box
|
|
return torch.from_numpy(aabb).float(), torch.from_numpy(transposed_point_cloud_transform).float()
|
|
|
|
|
|
@profiler.time_function
|
|
def load_friends_data(basedir, downscale_factor=1, split="train", include_semantics=True, include_point_cloud=False):
|
|
"""_summary_
|
|
|
|
Args:
|
|
basedir (_type_): _description_
|
|
downscale_factor (float, optional): _description_. Defaults to 1.0.
|
|
split (str, optional): _description_. Defaults to "train".
|
|
include_point_cloud (bool): whether or not to include the point cloud
|
|
"""
|
|
# pylint: disable=unused-argument
|
|
# pylint: disable=too-many-statements
|
|
|
|
# --- image filenames ----
|
|
images_data = read_images_binary(os.path.join(basedir, "colmap", "images.bin"))
|
|
# `image_path` is only the end of the filename, including the extension e.g., `.jpg`
|
|
image_paths = sorted(os.listdir(os.path.join(basedir, "images")))
|
|
|
|
image_path_to_image_id = {}
|
|
image_id_to_image_path = {}
|
|
for v in images_data.values():
|
|
image_path_to_image_id[v.name] = v.id
|
|
image_id_to_image_path[v.id] = v.name
|
|
# TODO: handle the splits differently
|
|
image_filenames = [os.path.join(basedir, "images", image_path) for image_path in image_paths]
|
|
|
|
# -- set the bounding box ---
|
|
aabb, transposed_point_cloud_transform = get_aabb_and_transform(basedir)
|
|
scene_bounds_original = SceneBounds(aabb=aabb)
|
|
# for shifting and rescale accoding to scene bounds
|
|
box_center = scene_bounds_original.get_center()
|
|
box_scale_factor = 5.0 / scene_bounds_original.get_diagonal_length() # the target diagonal length
|
|
scene_bounds = scene_bounds_original.get_centered_and_scaled_scene_bounds(box_scale_factor)
|
|
|
|
# --- intrinsics ---
|
|
cameras_data = read_cameras_binary(os.path.join(basedir, "colmap", "cameras.bin"))
|
|
intrinsics = []
|
|
for image_path in image_paths:
|
|
cam = cameras_data[image_path_to_image_id[image_path]]
|
|
assert len(cam.params) == 3
|
|
focal_length = cam.params[0] # f (fx and fy)
|
|
cx = cam.params[1] # cx
|
|
cy = cam.params[2] # cy
|
|
intrinsics.append([cx, cy, focal_length])
|
|
intrinsics = torch.tensor(intrinsics).float()
|
|
|
|
# --- camera_to_world (extrinsics) ---
|
|
camera_to_world = []
|
|
bottom_row = np.array([0, 0, 0, 1.0]).reshape(1, 4)
|
|
for image_path in image_paths:
|
|
image_data = images_data[image_path_to_image_id[image_path]]
|
|
rot = image_data.qvec2rotmat()
|
|
trans = image_data.tvec.reshape(3, 1)
|
|
c2w = np.concatenate([np.concatenate([rot, trans], 1), bottom_row], 0)
|
|
camera_to_world.append(c2w)
|
|
camera_to_world = torch.tensor(np.array(camera_to_world)).float()
|
|
camera_to_world = torch.inverse(camera_to_world)
|
|
camera_to_world[..., 1:3] *= -1
|
|
camera_to_world = transposed_point_cloud_transform @ camera_to_world
|
|
camera_to_world = camera_to_world[:, :3]
|
|
camera_to_world[..., 3] = (camera_to_world[..., 3] - box_center) * box_scale_factor # center and rescale
|
|
|
|
# --- semantics ---
|
|
semantics = Semantics()
|
|
if include_semantics:
|
|
thing_filenames = [
|
|
image_filename.replace("/images/", "/segmentations/thing/").replace(".jpg", ".png")
|
|
for image_filename in image_filenames
|
|
]
|
|
stuff_filenames = [
|
|
image_filename.replace("/images/", "/segmentations/stuff/").replace(".jpg", ".png")
|
|
for image_filename in image_filenames
|
|
]
|
|
panoptic_classes = load_from_json(os.path.join(basedir, "panoptic_classes.json"))
|
|
stuff_classes = panoptic_classes["stuff"]
|
|
stuff_colors = torch.tensor(panoptic_classes["stuff_colors"], dtype=torch.float32) / 255.0
|
|
thing_classes = panoptic_classes["thing"]
|
|
thing_colors = torch.tensor(panoptic_classes["thing_colors"], dtype=torch.float32) / 255.0
|
|
semantics = Semantics(
|
|
stuff_classes=stuff_classes,
|
|
stuff_colors=stuff_colors,
|
|
stuff_filenames=stuff_filenames,
|
|
thing_classes=thing_classes,
|
|
thing_colors=thing_colors,
|
|
thing_filenames=thing_filenames,
|
|
)
|
|
|
|
# Possibly include the sparse point cloud from COLMAP in the dataset inputs.
|
|
# NOTE(ethan): this will be common across the different splits.
|
|
point_cloud = PointCloud()
|
|
if include_point_cloud:
|
|
points_3d = read_pointsTD_binary(os.path.join(basedir, "colmap", "points3D.bin"))
|
|
xyz = torch.tensor(np.array([p_value.xyz for p_id, p_value in points_3d.items()])).float()
|
|
rgb = torch.tensor(np.array([p_value.rgb for p_id, p_value in points_3d.items()])).float()
|
|
xyz_h = torch.cat([xyz, torch.ones_like(xyz[..., :1])], -1)
|
|
xyz = (xyz_h @ transposed_point_cloud_transform.T)[..., :3]
|
|
xyz = (xyz - box_center) * box_scale_factor # center and rescale
|
|
point_cloud.xyz = xyz
|
|
point_cloud.rgb = rgb
|
|
|
|
dataset_inputs = DatasetInputs(
|
|
image_filenames=image_filenames,
|
|
downscale_factor=downscale_factor,
|
|
intrinsics=intrinsics / downscale_factor,
|
|
camera_to_world=camera_to_world,
|
|
semantics=semantics,
|
|
point_cloud=point_cloud,
|
|
scene_bounds=scene_bounds,
|
|
)
|
|
return dataset_inputs
|