diff --git a/GroundTruthAnnotations.md b/GroundTruthAnnotations.md index a926902be..e68a6de02 100644 --- a/GroundTruthAnnotations.md +++ b/GroundTruthAnnotations.md @@ -1,20 +1,28 @@ # Ground-Truth Annotations +### Changelog + +- 07/03/23: Add specification for Blender's built-in annotations. Save built-in annotations as numpy arrays. Add more information to Objects_XXXX_XX_XX.json. Significant changes to built-in segmentation masks (fixes & filenames). Improve visualizations for built-in annotations. Always save camera parameters in frames/. Update docs. + +### Agenda + +- Save forward and backward flow for both built-in and advanced annotations. +- Compute flow occlusion using forward-backward consistency. +- Export scene geometry in .ply format. + +**Want annotations that we don't currently support? [Fill out a request!](https://github.com/princeton-vl/infinigen/issues/new?assignees=&labels=&projects=&template=request.md&title=%5BREQUEST%5D)** + ## Default Annotations from Blender Infinigen can produce some dense annotations using Blender's built-in render passes. Users may prefer to use these annotations over our extended annotation system's since it requires only the bare-minimum installation. It is also able to run without a GPU. These annotations are produced when using the `--pipeline_configs blender_gt` ground truth extraction config in [manage_datagen_jobs.py](/README.md#generate-images-in-one-command), or can be done manually as shown in the final step of the [Hello-World](/README.md#generate-a-scene-step-by-step) example. -### Specification +## Advanced Annotation Pipeline :large_blue_diamond: -**Coming Soon** +We also provide a separate pipeline for extracting the full set of annotations from each image or scene. Features only supported using this annotation method will be denoted with :large_blue_diamond:. -## OpenGL-Based Annotation Pipeline - -We also provide a separate pipeline for extracting the full set of annotations from each image or scene. - -This section will allow you to use our own `--pipeline_configs opengl_gt` ground truth extraction config, which provides additional labels such as occlusion boundaries, sub-object segmentation, 3D flow and easy 3D bounding boxes. If you do not need these features, we recommend using the [default annotations](#default-annotations-from-blender). This section is intended for computer vision researchers and power-users. +This will allow you to use our own `--pipeline_configs opengl_gt` ground truth extraction config, which provides additional labels such as occlusion boundaries, sub-object segmentation, 3D flow and easy 3D bounding boxes. If you do not need these features, we recommend using the [default annotations](#default-annotations-from-blender). This section is intended for computer vision researchers and power-users. ### Installation @@ -77,7 +85,7 @@ python tools/summarize.py outputs/helloworld # creating outputs/helloworld/summa python tools/ground_truth/segmentation_lookup.py outputs/helloworld 1 --query cactus ``` -### Specification +## Specification **File structure:** @@ -103,7 +111,7 @@ The resulting `/summary.json` will contains all file paths in the `` and `` are typically both "00" in the monocular setting; `` is typically "npy" or "png" for the the actual data and the visualization, respectively; `` is a 0-padded 4-digit number, e.g. "0013". `` can be "SurfaceNormal", "Depth", etc. For example `summary_json["SurfaceNormal"]["npy"]["00"]["00"]["0001"]` -> `'frames/SurfaceNormal_0001_00_00.npy'` -*Note: Currently our ground-truth has only been tested for the aspect-ratio 16-9.* +*Note: Currently our advanced ground-truth has only been tested for the aspect-ratio 16-9.* **Depth** @@ -126,6 +134,8 @@ python tools/ground_truth/rigid_warp.py Surface Normals are stored as a 1080 x 1920 x 3 32-bit floating point numpy array. +The coordinate system for the surface normals is +X -> Right, +Y -> Up, +Z Backward. + *Path:* `summary_json["SurfaceNormal"]["npy"]["00"]["00"]["0001"]` -> `frames/SurfaceNormal_0001_00_00.npy` *Visualization:* `summary_json["SurfaceNormal"]["png"]["00"]["00"]["0001"]` -> `frames/SurfaceNormal_0001_00_00.png` @@ -134,7 +144,7 @@ Surface Normals are stored as a 1080 x 1920 x 3 32-bit floating point numpy arra

-**Occlusion Boundaries** +### Occlusion Boundaries :large_blue_diamond: Occlusion Boundaries are stored as a 2160 x 3840 png, with 255 indicating a boundary and 0 otherwise. @@ -144,7 +154,7 @@ Occlusion Boundaries are stored as a 2160 x 3840 png, with 255 indicating a boun

-**Optical Flow / Scene Flow** +### Optical Flow Optical Flow / Scene Flow is stored as a 2160 x 3840 x 3 32-bit floating point numpy array. @@ -152,6 +162,8 @@ Optical Flow / Scene Flow is stored as a 2160 x 3840 x 3 32-bit floating point n Channels 1 & 2 are standard optical flow. Note that the units of optical flow are in pixels measured in the resolution of the *original image*. So if the rendered image is 1080 x 1920, you would want to average-pool this array by 2x. +**3D Motion Vectors** :large_blue_diamond: + Channel 3 is the depth change between this frame and the next. To see an example of how optical flow can be used to warp one frame to the next, run @@ -162,9 +174,11 @@ python tools/ground_truth/optical_flow_warp.py *Path:* `summary_json["Flow3D"]["npy"]["00"]["00"]["0001"]` -> `frames/Flow3D_0001_00_00.npy` -*Visualization:* `summary_json["Flow3D"]["png"]["00"]["00"]["0001"]` -> `frames/ObjectSegmentation_0001_00_00.png` +*Visualization:* `summary_json["Flow3D"]["png"]["00"]["00"]["0001"]` -> `frames/Flow3D_0001_00_00.png` + +For the built-in versions, replace `Flow3D` with `Flow`. -**Optical Flow Occlusion** +### Optical Flow Occlusion :large_blue_diamond: The mask of occluded pixels for the aforementioned optical flow is stored as a 2160 x 3840 png, with 255 indicating a co-visible pixel and 0 otherwise. @@ -172,51 +186,58 @@ The mask of occluded pixels for the aforementioned optical flow is stored as a 2 *Path/Visualization:* `summary_json["Flow3DMask"]["png"]["00"]["00"]["0001"]` -> `frames/Flow3DMask_0001_00_00.png` -**Camera Intrinsics** +### Camera Intrinsics Infinigen renders images using a pinhole camera model. The resulting camera intrinsics for each frame are stored as a 3 x 3 numpy matrix. -*Path:* `summary_json["Camera Intrinsics"]["npy"]["00"]["00"]["0001"]` -> `saved_mesh/frame_0001/cameras/K_0001_00_00.npy` +*Path:* `summary_json["Camera Intrinsics"]["npy"]["00"]["00"]["0001"]` -> `frames/K_0001_00_00.npy` -**Camera Extrinsics** +### Camera Extrinsics The camera pose is stored as a 4 x 4 numpy matrix mapping from camera coordinates to world coordinates. As is standard in computer vision, the assumed world coordinate system in the saved camera poses is +X -> Right, +Y -> Down, +Z Forward. This is opposed to how Blender internally represents geometry, with flipped Y and Z axes. -*Path:* `summary_json["Camera Pose"]["npy"]["00"]["00"]["0001"]` -> `saved_mesh/frame_0001/cameras/T_0001_00_00.npy` +*Path:* `summary_json["Camera Pose"]["npy"]["00"]["00"]["0001"]` -> `frames/T_0001_00_00.npy` -**Panoptic Segmentation and 3D Bounding Boxes** +### Panoptic Segmentation Infinigen saves three types of semantic segmentation masks: 1) Object Segmentation 2) Tag Segmentation 3) Instance Segmentation -*Object Segmentation* distinguishes individual blender objects, and is stored as a 2160 x 3840 32-bit integer numpy array. The association between each integer in the mask and the related object is stored in Objects_XXXX_XX_XX.json. The definition of "object" is imposed by Blender; generally large or complex assets such as the terrain, trees, or animals are considered one singular object, while a large number of smaller assets (e.g. grass, coral) may be grouped together if they are using instanced-geometry for their implementation. +*Object Segmentation* distinguishes individual blender objects, and is stored as a 2160 x 3840 32-bit integer numpy array. Each integer in the mask maps to an object in Objects_XXXX_XX_XX.json with the same value for the `"object_index"` field. The definition of "object" is imposed by Blender; generally large or complex assets such as the terrain, trees, or animals are considered one singular object, while a large number of smaller assets (e.g. grass, coral) may be grouped together if they are using instanced-geometry for their implementation. -*Tag Segmentation* distinguishes vertices based on their semantic tags, and is stored as a 2160 x 3840 64-bit integer numpy array. Infinigen tags all vertices with an integer which can be associated to a list of semantic labels in `MaskTag.json`. Compared to Object Segmentation, Infinigen's tagging system is less automatic but much more flexible. Missing features in the tagging system are usually possible and straightforward to implement, wheras in the automaically generated Object Segmentation they are not. +*Instance Segmentation* distinguishes individual instances of a single object from one another (e.g. separate blades of grass, separate ferns, etc.), and is stored as a 2160 x 3840 32-bit integer numpy array. Each integer in this mask is the *instance-id* for a particular instance, which is unique for that object as defined in the Object Segmentation mask and Objects_XXXX_XX_XX.json. -*Instance Segmentation* distinguishes individual instances of a single object from one another (e.g. separate blades of grass, separate ferns, etc.), and is stored as a 2160 x 3840 32-bit integer numpy array. Each integer in this mask is the *instance-id* for a particular instance, which is unique for that object as defined in the Object Segmentation mask and Objects_XXXX_XX_XX.json. The list of **3D bounding boxes** for each instance are also defined in the `Objects_XXXX_XX_XX.json`. *Paths:* `summary_json["ObjectSegmentation"]["npy"]["00"]["00"]["0001"]` -> `frames/ObjectSegmentation_0001_00_00.npy` -`summary_json["TagSegmentation"]["npy"]["00"]["00"]["0001"]` -> `frames/TagSegmentation_0001_00_00.npy` - `summary_json["InstanceSegmentation"]["npy"]["00"]["00"]["0001"]` -> `frames/InstanceSegmentation_0001_00_00.npy` `summary_json["Objects"]["json"]["00"]["00"]["0001"]` -> `frames/Objects_0001_00_00.json` -`summary_json["Mask Tags"][]` -> `fine/MaskTag.json` - *Visualizations:* `summary_json["ObjectSegmentation"]["png"]["00"]["00"]["0001"]` -> `frames/ObjectSegmentation_0001_00_00.png` -`summary_json["TagSegmentation"]["png"]["00"]["00"]["0001"]` -> `frames/TagSegmentation_0001_00_00.png` - `summary_json["InstanceSegmentation"]["png"]["00"]["00"]["0001"]` -> `frames/InstanceSegmentation_0001_00_00.png` -Generally, most useful panoptic segmentation masks can be constructed by combining the aforementioned three arrays in some way. As an example, to visualize the 2D and 3D bounding boxes for objects with the *blender_rock* semantic tag in the hello world scene, run +#### **Tag Segmentation** :large_blue_diamond: + +*Tag Segmentation* distinguishes vertices based on their semantic tags, and is stored as a 2160 x 3840 64-bit integer numpy array. Infinigen tags all vertices with an integer which can be associated to a list of semantic labels in `MaskTag.json`. Compared to Object Segmentation, Infinigen's tagging system is less automatic but much more flexible. Missing features in the tagging system are usually possible and straightforward to implement, wheras in the automaically generated Object Segmentation they are not. + +*Paths:* + +`summary_json["TagSegmentation"]["npy"]["00"]["00"]["0001"]` -> `frames/TagSegmentation_0001_00_00.npy` + +`summary_json["Mask Tags"][]` -> `fine/MaskTag.json` + +*Visualization:* + +`summary_json["TagSegmentation"]["png"]["00"]["00"]["0001"]` -> `frames/TagSegmentation_0001_00_00.png` + +Generally, most useful panoptic segmentation masks can be constructed by combining the aforementioned three arrays in some way. As an example, to visualize the 2D and [3D bounding boxes](#object-metadata-and-3d-bounding-boxes) for objects with the *blender_rock* semantic tag in the hello world scene, run ``` python tools/ground_truth/segmentation_lookup.py outputs/helloworld 1 --query blender_rock --boxes python tools/ground_truth/bounding_boxes_3d.py outputs/helloworld 1 --query blender_rock @@ -247,3 +268,32 @@ python tools/ground_truth/segmentation_lookup.py outputs/helloworld 1 --query wa

+ +### Object Metadata and 3D bounding boxes + +Each item in `Objects_0001_00_00.json` also contains other metadata about each object: +``` +# Load object meta data +object_data = json.loads((Path("output/helloworld") / summary_json["Objects"]["json"]["00"]["00"]["0001"]).read_text()) + +# select nth object +obj = object_data[n] + +obj["children"] # list of object indices for children +obj["object_index"] # object index, for lookup in the Object Segmentation mask +obj["num_verts"] # number of vertices +obj["num_faces"] # number of faces (n-gons, not triangles) +obj["name"] # obvious +obj["unapplied_modifiers"] # names of unapplied blender modifiers +obj["materials"] # materials used +``` + +More fields :large_blue_diamond: +``` +obj["tags"] # list of tags which appear on at least one vertex +obj["min"] # min-corner of bounding box, in object coordinates +obj["max"] # max-corner of bounding box, in object coordinates +obj["model_matrices"] # mapping from instance-ids to 4x4 obj->world transformation matrices +``` + +The **3D bounding box** for each instance can be computed using `obj["min"]`, `obj["max"]`, `obj["model_matrices"]`. For an example, refer to [the bounding_boxes_3d.py example above](#tag-segmentation-large_blue_diamond). diff --git a/README.md b/README.md index e95ef886d..89bbf0ec1 100644 --- a/README.md +++ b/README.md @@ -97,11 +97,13 @@ Install [WSL2](https://infinigen.org/docs/installation/intro#setup-for-windows) :warning: **Known issue** : We are actively fixing an issue which causes commands not to be reproducible on many platforms. The same command may produce multiple rearranged scenes with different runtimes and memory requirements.

- - + + + +

-This guide will show you how to generate an image and it's corresponding depth ground-truth, similar to those shown above. +This guide will show you how to generate an image and it's corresponding ground-truth, similar to those shown above. #### Generate a scene step by step Infinigen generates scenes by running multiple tasks (usually executed automatically, like in [Generate image(s) in one command](#generate-images-in-one-command)). Here we will run them one by one to demonstrate. These commands take approximately 10 minutes and 16GB of memory to execute on an M1 Mac or Linux Desktop. @@ -125,6 +127,10 @@ $BLENDER -noaudio --background --python generate.py -- --seed 0 --task render -g Output logs should indicate what the code is working on. Use `--debug` for even more detail. After each command completes you can inspect it's `--output_folder` for results, including running `$BLENDER outputs/helloworld/coarse/scene.blend` or similar to view blender files. We hide many meshes by default for viewport stability; to view them, click "Render" or use the UI to unhide them. +#### [Extended ground-truth & docmentation](./GroundTruthAnnotations.md) + +We also provide a (optional) separate pipeline for extracting the full set of annotations from each image or scene. Refer to [GroundTruthAnnotations.md](./GroundTruthAnnotations.md) for compilation instructions, data format specifications and an extended "Hello World". + #### Generate image(s) in one command We provide `tools/manage_datagen_jobs.py`, a utility which runs these or similar steps automatically. diff --git a/images/InstanceSegmentation_0001_00_00.png b/images/InstanceSegmentation_0001_00_00.png new file mode 100644 index 000000000..7d2c0b995 Binary files /dev/null and b/images/InstanceSegmentation_0001_00_00.png differ diff --git a/images/SurfaceNormal_0001_00_00.png b/images/SurfaceNormal_0001_00_00.png new file mode 100644 index 000000000..3f463129a Binary files /dev/null and b/images/SurfaceNormal_0001_00_00.png differ diff --git a/process_mesh/blender_object.cpp b/process_mesh/blender_object.cpp index efdfe0506..b06d58056 100644 --- a/process_mesh/blender_object.cpp +++ b/process_mesh/blender_object.cpp @@ -73,10 +73,15 @@ json BaseBlenderObject::compute_bbox(const std::vector &indices, c const std::set unique_tags(tag_lookup.begin(), tag_lookup.end()); json output = { - {"model matrices", json_serializable_model_matrices}, + {"model_matrices", json_serializable_model_matrices}, {"tags", std::vector(unique_tags.begin(), unique_tags.end())}, - {"name", name}, - {"object index", obj_index} + {"name", info.name}, + {"num_verts", info.num_verts}, + {"num_faces", info.num_faces}, + {"children", info.children}, + {"materials", info.materials}, + {"unapplied_modifiers", info.unapplied_modifiers}, + {"object_index", info.index} }; if ((num_verts > 0) && ((max - min).norm() > 1e-4)){ @@ -91,7 +96,7 @@ json BaseBlenderObject::compute_bbox(const std::vector &indices, c } BaseBlenderObject::BaseBlenderObject(const BufferArrays ¤t_buf, const BufferArrays &next_buf, const std::vector &instance_ids, const ObjectInfo& object_info, const ObjectType tp, int attrib_stride) - : num_verts(current_buf.indices.size()), type(tp), name(object_info.name), num_instances(instance_ids.size()), obj_index(object_info.index) { + : num_verts(current_buf.indices.size()), type(tp), info(object_info), num_instances(instance_ids.size()) { const std::vector &model_matrices = current_buf.get_instances(instance_ids); const std::vector &model_matrices_next = next_buf.get_instances(instance_ids); @@ -172,7 +177,7 @@ MeshBlenderObject::~MeshBlenderObject(){} void MeshBlenderObject::draw(Shader &shader) const { const auto t1 = std::chrono::high_resolution_clock::now(); - shader.setInt("object_index", obj_index); + shader.setInt("object_index", info.index); glBindVertexArray(VAO); glDrawElementsInstanced(GL_LINES_ADJACENCY, num_verts, GL_UNSIGNED_INT, 0, num_instances); glCheckError(); diff --git a/process_mesh/blender_object.hpp b/process_mesh/blender_object.hpp index 4f48977b4..49fd27027 100644 --- a/process_mesh/blender_object.hpp +++ b/process_mesh/blender_object.hpp @@ -17,18 +17,25 @@ using json = nlohmann::json; struct ObjectInfo { - int index, num_instances; + int index, num_instances, num_faces, num_verts; std::string name, type, mesh_id, npz_filename; + std::vector children; + std::vector materials, unapplied_modifiers; ObjectInfo(){} - ObjectInfo(nlohmann::json_abi_v3_11_2::json instance_item) : - name(instance_item["object_name"].get()), - type(instance_item["object_type"].get()), - index(instance_item["object_idx"].get()), - num_instances(instance_item["num_instances"].get()), - mesh_id(instance_item["mesh_id"].get()), - npz_filename(instance_item["filename"].get()){} + ObjectInfo(const json instance_item) : + name(instance_item["object_name"]), + type(instance_item["object_type"]), + index(instance_item["object_idx"]), + num_instances(instance_item["num_instances"]), + num_faces(instance_item["num_faces"]), + num_verts(instance_item["num_verts"]), + children(instance_item["children"]), + materials(instance_item["materials"]), + unapplied_modifiers(instance_item["unapplied_modifiers"]), + mesh_id(instance_item["mesh_id"]), + npz_filename(instance_item["filename"]){} }; @@ -54,8 +61,7 @@ class BaseBlenderObject json bounding_box; const ObjectType type; - const std::string name; - const int obj_index; + const ObjectInfo info; BaseBlenderObject(const BufferArrays ¤t_buf, const BufferArrays &next_buf, const std::vector &instance_ids, const ObjectInfo& object_info, const ObjectType tp, int attrib_stride); virtual ~BaseBlenderObject(); diff --git a/process_mesh/main.cpp b/process_mesh/main.cpp index 12333b62e..4c7e214b7 100644 --- a/process_mesh/main.cpp +++ b/process_mesh/main.cpp @@ -32,7 +32,7 @@ #include "utils.hpp" #include "io.hpp" -#define VERSION "1.34" +#define VERSION "1.35" using std::cout, std::cerr, std::endl; diff --git a/requirements.txt b/requirements.txt index 4ed78b358..ff2d23c06 100644 --- a/requirements.txt +++ b/requirements.txt @@ -15,7 +15,8 @@ trimesh einops mesh_to_sdf geomdl -numpy==1.21.5 +numpy>=1.22 +numba wandb zarr jinja2 diff --git a/worldgen/config/base.gin b/worldgen/config/base.gin index 45c00ee7b..24e1772c5 100644 --- a/worldgen/config/base.gin +++ b/worldgen/config/base.gin @@ -103,8 +103,8 @@ full/render_image.num_samples = 8192 render_image.adaptive_threshold = 0.005 full/render_image.flat_shading = False -flat/render_image.min_samples = 10 -flat/render_image.num_samples = 10 +flat/render_image.min_samples = 1 +flat/render_image.num_samples = 1 flat/render_image.flat_shading = True full/render_image.passes_to_save = [ ['diffuse_direct', 'DiffDir'], diff --git a/worldgen/core.py b/worldgen/core.py index 7c3b7068c..5a075b1d6 100644 --- a/worldgen/core.py +++ b/worldgen/core.py @@ -202,8 +202,8 @@ def save_meshes(scene_seed, output_folder, frame_range, camera_id, resample_idx= frame_info_folder = Path(output_folder) / f"frame_{frame_idx:04d}" frame_info_folder.mkdir(parents=True, exist_ok=True) logging.info(f"Working on frame {frame_idx}") - exporting.save_obj_and_instances(str(frame_info_folder / "mesh"), previous_frame_mesh_id_mapping, current_frame_mesh_id_mapping) - cam_util.save_camera_parameters(camera_rig_id, [subcam_id], str(frame_info_folder / "cameras"), frame_idx) + exporting.save_obj_and_instances(frame_info_folder / "mesh", previous_frame_mesh_id_mapping, current_frame_mesh_id_mapping) + cam_util.save_camera_parameters(camera_rig_id, [subcam_id], frame_info_folder / "cameras", frame_idx) previous_frame_mesh_id_mapping = frozendict(current_frame_mesh_id_mapping) current_frame_mesh_id_mapping.clear() diff --git a/worldgen/placement/camera.py b/worldgen/placement/camera.py index c44ce1500..45b9fccee 100644 --- a/worldgen/placement/camera.py +++ b/worldgen/placement/camera.py @@ -26,7 +26,7 @@ from mathutils.bvhtree import BVHTree from assets.utils.decorate import toggle_hide -from rendering.post_render import depth_to_jet +from rendering.post_render import colorize_depth from tqdm import tqdm, trange from placement import placement @@ -532,5 +532,5 @@ def get_camera_trajectory(): assert dist > dist_diff depth_output[H-y-1,x] = dist - dist_diff - color_depth = depth_to_jet(depth_output) + color_depth = colorize_depth(depth_output) imageio.imwrite(f"color_depth.png", color_depth) diff --git a/worldgen/rendering/post_render.py b/worldgen/rendering/post_render.py index a616c5c48..29f7c7ef2 100644 --- a/worldgen/rendering/post_render.py +++ b/worldgen/rendering/post_render.py @@ -4,8 +4,6 @@ # Authors: Lahav Lipson -import time -import warnings import argparse import os os.environ["OPENCV_IO_ENABLE_OPENEXR"]="1" # This must be done BEFORE import cv2. @@ -17,14 +15,30 @@ from pathlib import Path from imageio import imwrite -def flow_to_colorwheel(flow_path): - assert flow_path.exists() and flow_path.suffix == ".exr" - optical_flow = cv2.imread(str(flow_path), cv2.IMREAD_ANYCOLOR | cv2.IMREAD_ANYDEPTH) +def load_exr(path): + assert Path(path).exists() and Path(path).suffix == ".exr" + return cv2.imread(str(path), cv2.IMREAD_ANYCOLOR | cv2.IMREAD_ANYDEPTH) + +load_flow = load_exr +load_depth = lambda p: load_exr(p)[..., 0] +load_normals = lambda p: load_exr(p)[...,[2,0,1]] * np.array([-1.,1.,1.]) +load_seg_mask = lambda p: load_exr(p)[...,2].astype(np.int64) +load_uniq_inst = lambda p: load_exr(p).view(np.int32) + +def colorize_flow(optical_flow): flow_uv = optical_flow[...,:2] flow_color = flow_vis.flow_to_color(flow_uv, convert_to_bgr=False) return flow_color -def depth_to_jet(depth, scale_vmin=1.0): +def colorize_normals(surface_normals): + assert surface_normals.max() < 1+1e-4 + assert surface_normals.min() > -1-1e-4 + norm = np.linalg.norm(surface_normals, axis=2) + color = np.round((surface_normals + 1) * (255/2)).astype(np.uint8) + color[norm < 1e-4] = 0 + return color + +def colorize_depth(depth, scale_vmin=1.0): valid = (depth > 1e-3) & (depth < 1e4) vmin = depth[valid].min() * scale_vmin vmax = depth[valid].max() @@ -34,46 +48,49 @@ def depth_to_jet(depth, scale_vmin=1.0): depth[~valid] = 1 return np.ascontiguousarray(depth[...,:3] * 255, dtype=np.uint8) -def mask_to_color(mask_path, color_seed=None): - assert mask_path.exists() and mask_path.suffix == ".exr" - exr = cv2.imread(str(mask_path), cv2.IMREAD_ANYCOLOR | cv2.IMREAD_ANYDEPTH)[...,2].astype(np.int64) - H,W = exr.shape - top = exr.max()+1 - if color_seed is None: - color_seed = int(time.time()*1000)%1000 - perm = np.random.RandomState(color_seed).permutation(top) - output_image = np.zeros((H, W, 3), dtype=np.uint8) - for i in range(top): - clr = (np.asarray(colorsys.hsv_to_rgb(i / top, 0.9, 0.8))*255).astype(np.uint8) - output_image[exr == perm[i]] = clr - return output_image - -def exr_depth_to_jet(depth_path, scale_vmin=1.0): - assert depth_path.exists() and depth_path.suffix == ".exr" - depth = cv2.imread(str(depth_path), cv2.IMREAD_ANYCOLOR | cv2.IMREAD_ANYDEPTH)[..., 0] - return depth_to_jet(depth) +def colorize_int_array(data, color_seed=0): + H, W, *_ = data.shape + data = data.reshape((H * W, -1)) + uniq, indices = np.unique(data, return_inverse=True, axis=0) + random_states = [np.random.RandomState(e[:2].astype(np.uint32) + color_seed) for e in uniq] + unique_colors = (np.asarray([colorsys.hsv_to_rgb(s.uniform(0, 1), s.uniform(0.1, 1), 1) for s in random_states]) * 255).astype(np.uint8) + return unique_colors[indices].reshape((H, W, 3)) if __name__ == "__main__": parser = argparse.ArgumentParser() parser.add_argument('--flow_path', type=Path, default=None) parser.add_argument('--depth_path', type=Path, default=None) parser.add_argument('--seg_path', type=Path, default=None) + parser.add_argument('--uniq_inst_path', type=Path, default=None) + parser.add_argument('--normals_path', type=Path, default=None) args = parser.parse_args() if args.flow_path is not None: - flow_color = flow_to_colorwheel(args.flow_path) + flow_color = colorize_flow(load_flow(args.flow_path)) output_path = args.flow_path.with_suffix('.png') imwrite(output_path, flow_color) print(f"Wrote {output_path}") + if args.normals_path is not None: + normal_color = colorize_normals(load_normals(args.normals_path)) + output_path = args.normals_path.with_suffix('.png') + imwrite(output_path, normal_color) + print(f"Wrote {output_path}") + if args.depth_path is not None: - depth_color = exr_depth_to_jet(args.depth_path) + depth_color = colorize_depth(load_depth(args.depth_path)) output_path = args.depth_path.with_suffix('.png') imwrite(output_path, depth_color) print(f"Wrote {output_path}") + if args.uniq_inst_path is not None: + mask_color = colorize_int_array(load_uniq_inst(args.uniq_inst_path)) + output_path = args.uniq_inst_path.with_suffix('.png') + imwrite(output_path, mask_color) + print(f"Wrote {output_path}") + if args.seg_path is not None: - mask_color = mask_to_color(args.seg_path) + mask_color = colorize_int_array(load_seg_mask(args.seg_path)) output_path = args.seg_path.with_suffix('.png') imwrite(output_path, mask_color) print(f"Wrote {output_path}") diff --git a/worldgen/rendering/render.py b/worldgen/rendering/render.py index c1a714911..ae5e11dbd 100644 --- a/worldgen/rendering/render.py +++ b/worldgen/rendering/render.py @@ -7,30 +7,31 @@ # - Hei Law - Initial version +import json +import logging import os import time import warnings + import bpy import gin -import json -import os -import cv2 import numpy as np -import logging +from imageio import imread, imwrite + from infinigen_gpl.extras.enable_gpu import enable_gpu -from imageio import imwrite, imread -from pathlib import Path +from nodes.node_wrangler import Nodes, NodeWrangler from placement import camera as cam_util -from nodes.node_wrangler import NodeWrangler, Nodes -from rendering.post_render import exr_depth_to_jet, flow_to_colorwheel, mask_to_color +from rendering.post_render import (colorize_depth, colorize_flow, + colorize_normals, colorize_int_array, + load_depth, load_flow, load_normals, + load_seg_mask, load_uniq_inst) +from surfaces import surface +from util import blender as butil +from util import exporting as exputil +from util.logging import Timer from .auto_exposure import nodegroup_auto_exposure -from util.camera import get_calibration_matrix_K_from_blender -from util.logging import Timer -from surfaces import surface -from util import blender as butil, exporting as exputil - TRANSPARENT_SHADERS = {Nodes.TranslucentBSDF, Nodes.TransparentBSDF} logger = logging.getLogger(__name__) @@ -54,30 +55,31 @@ def remove_translucency(): assert not fac_soc.is_linked fac_soc.default_value = 0.0 -def save_and_set_pass_indices(output_folder): - file_tree = {} - set_pass_indices(bpy.context.scene.collection, 1, file_tree) - json_object = json.dumps(file_tree) - (output_folder / "object_tree.json").write_text(json_object) - -def set_pass_indices(parent_collection, index, tree_output): - for child_obj in parent_collection.objects: - child_obj.pass_index = index +def set_pass_indices(): + tree_output = {} + index = 1 + for obj in bpy.data.objects: + if obj.hide_render: + continue + if obj.pass_index == 0: + obj.pass_index = index + index += 1 object_dict = { - "type": child_obj.type, "pass_index": index, - "bbox": np.asarray(child_obj.bound_box[:]).tolist(), - "matrix_world": np.asarray(child_obj.matrix_world[:]).tolist(), + "type": obj.type, "object_index": obj.pass_index, "children": [] } - if child_obj.type == "MESH": - object_dict['polycount'] = len(child_obj.data.polygons) - object_dict['materials'] = child_obj.material_slots.keys() - object_dict['unapplied_modifiers'] = child_obj.modifiers.keys() - tree_output[child_obj.name] = object_dict + if obj.type == "MESH": + object_dict['num_verts'] = len(obj.data.vertices) + object_dict['num_faces'] = len(obj.data.polygons) + object_dict['materials'] = obj.material_slots.keys() + object_dict['unapplied_modifiers'] = obj.modifiers.keys() + tree_output[obj.name] = object_dict + for child_obj in obj.children: + if child_obj.pass_index == 0: + child_obj.pass_index = index + index += 1 + object_dict["children"].append(child_obj.pass_index) index += 1 - for col in parent_collection.children: - tree_output[col.name] = {"type": "Collection", "hide_viewport": col.hide_viewport, "children": {}} - index = set_pass_indices(col, index, tree_output=tree_output[col.name]["children"]) - return index + return tree_output # Can be pasted directly into the blender console def make_clay(): @@ -141,7 +143,7 @@ def configure_compositor_output(nw, frames_folder, image_denoised, image_noisy, image = image_denoised if use_denoised else image_noisy nw.links.new(image, file_output_node.inputs['Image']) if saving_ground_truth: - slot_input.path = 'Unique_Instances' + slot_input.path = 'UniqueInstances' else: image_exr_output_node = nw.new_node(Nodes.OutputFile, attrs={ "base_path": str(frames_folder), @@ -166,10 +168,6 @@ def shader_random(nw: NodeWrangler): nw.new_node(Nodes.MaterialOutput, input_kwargs={'Surface': white_noise_texture.outputs["Color"]}) -def apply_random(obj, selection=None, **kwargs): - surface.add_material(obj, shader_random, selection=selection) - - def global_flat_shading(): for obj in bpy.context.scene.view_layers['ViewLayer'].objects: @@ -195,7 +193,14 @@ def global_flat_shading(): bpy.ops.object.material_slot_remove() for obj in bpy.context.scene.view_layers['ViewLayer'].objects: - apply_random(obj) + surface.add_material(obj, shader_random) + for mat in bpy.data.materials: + nw = NodeWrangler(mat.node_tree) + shader_random(nw) + + nw = NodeWrangler(bpy.data.worlds["World"].node_tree) + for link in nw.links: + nw.links.remove(link) @gin.configurable def render_image( @@ -249,7 +254,10 @@ def render_image( if flat_shading: with Timer("Set object indices"): - save_and_set_pass_indices(frames_folder) + object_data = set_pass_indices() + json_object = json.dumps(object_data, indent=4) + first_frame = bpy.context.scene.frame_start + (frames_folder / f"Objects_{first_frame:04d}_{camera_rig_id:02d}_{subcam_id:02d}.json").write_text(json_object) with Timer("Flat Shading"): global_flat_shading() @@ -295,40 +303,50 @@ def render_image( with Timer(f"Actual rendering"): bpy.ops.render.render(animation=True) - if flat_shading: - with Timer(f"Post Processing"): - for frame in range(bpy.context.scene.frame_start, bpy.context.scene.frame_end + 1): + with Timer(f"Post Processing"): + for frame in range(bpy.context.scene.frame_start, bpy.context.scene.frame_end + 1): + if flat_shading: bpy.context.scene.frame_set(frame) - K = get_calibration_matrix_K_from_blender(camera.data) - cameras_folder = frames_folder / "cameras" - cameras_folder.mkdir(exist_ok=True, parents=True) - np.save( - frames_folder / f"K{frame:04d}_{camera_rig_id:02d}_{subcam_id:02d}.npy", - np.asarray(K, dtype=np.float64), - ) - np.save( - frames_folder / f"T{frame:04d}_{camera_rig_id:02d}_{subcam_id:02d}.npy", - np.asarray(camera.matrix_world, dtype=np.float64), - ) - - # Save flow visualization. Takes about 3 seconds - flow_dst_path = frames_folder / f"Vector_{frame:04d}_{camera_rig_id:02d}_{subcam_id:02d}.exr" - if flow_dst_path.exists(): - flow_color = flow_to_colorwheel(flow_dst_path) - imwrite(flow_dst_path.with_name(f"Flow_{frame:04d}_{camera_rig_id:02d}_{subcam_id:02d}.png"), flow_color) - - # Save depth visualization. Also takes about 3 seconds - depth_dst_path = frames_folder / f"Depth_{frame:04d}_{camera_rig_id:02d}_{subcam_id:02d}.exr" - if depth_dst_path.exists(): - depth_color = exr_depth_to_jet(depth_dst_path) - imwrite(depth_dst_path.with_name(f"Depth_{frame:04d}_{camera_rig_id:02d}_{subcam_id:02d}.png"), depth_color) - - # Save Segmentation visualization. Also takes about 3 seconds - seg_dst_path = frames_folder / f"IndexOB_{frame:04d}_{camera_rig_id:02d}_{subcam_id:02d}.exr" - if seg_dst_path.exists(): - seg_color = mask_to_color(seg_dst_path) - imwrite(seg_dst_path.with_name(f"Segmentation_{frame:04d}_{camera_rig_id:02d}_{subcam_id:02d}.png"), seg_color) + output_stem = f"{frame:04d}_{camera_rig_id:02d}_{subcam_id:02d}" + + # Save flow visualization + flow_dst_path = frames_folder / f"Vector_{output_stem}.exr" + flow_array = load_flow(flow_dst_path) + np.save(flow_dst_path.with_name(f"Flow_{output_stem}.npy"), flow_array) + imwrite(flow_dst_path.with_name(f"Flow_{output_stem}.png"), colorize_flow(flow_array)) + flow_dst_path.unlink() + + # Save surface normal visualization + normal_dst_path = frames_folder / f"Normal_{output_stem}.exr" + normal_array = load_normals(normal_dst_path) + np.save(flow_dst_path.with_name(f"SurfaceNormal_{output_stem}.npy"), normal_array) + imwrite(flow_dst_path.with_name(f"SurfaceNormal_{output_stem}.png"), colorize_normals(normal_array)) + normal_dst_path.unlink() + + # Save depth visualization + depth_dst_path = frames_folder / f"Depth_{output_stem}.exr" + depth_array = load_depth(depth_dst_path) + np.save(flow_dst_path.with_name(f"Depth_{output_stem}.npy"), depth_array) + imwrite(depth_dst_path.with_name(f"Depth_{output_stem}.png"), colorize_depth(depth_array)) + depth_dst_path.unlink() + + # Save segmentation visualization + seg_dst_path = frames_folder / f"IndexOB_{output_stem}.exr" + seg_mask_array = load_seg_mask(seg_dst_path) + np.save(flow_dst_path.with_name(f"ObjectSegmentation_{output_stem}.npy"), seg_mask_array) + imwrite(seg_dst_path.with_name(f"ObjectSegmentation_{output_stem}.png"), colorize_int_array(seg_mask_array)) + seg_dst_path.unlink() + + # Save unique instances visualization + uniq_inst_path = frames_folder / f"UniqueInstances_{output_stem}.exr" + uniq_inst_array = load_uniq_inst(uniq_inst_path) + np.save(flow_dst_path.with_name(f"InstanceSegmentation_{output_stem}.npy"), uniq_inst_array) + imwrite(uniq_inst_path.with_name(f"InstanceSegmentation_{output_stem}.png"), colorize_int_array(uniq_inst_array)) + uniq_inst_path.unlink() + else: + cam_util.save_camera_parameters(camera_rig_id, [subcam_id], frames_folder, frame) + for file in tmp_dir.glob('*.png'): file.unlink() diff --git a/worldgen/tools/ground_truth/bounding_boxes_3d.py b/worldgen/tools/ground_truth/bounding_boxes_3d.py index 065fbea93..c9aeab600 100644 --- a/worldgen/tools/ground_truth/bounding_boxes_3d.py +++ b/worldgen/tools/ground_truth/bounding_boxes_3d.py @@ -48,9 +48,6 @@ def calc_bbox_pts(min_pt, max_pt): if __name__ == "__main__": - import os,sys - sys.path.append(os.getcwd()) - parser = argparse.ArgumentParser() parser.add_argument('folder', type=Path) parser.add_argument('frame', type=int) diff --git a/worldgen/tools/ground_truth/depth_to_normals.py b/worldgen/tools/ground_truth/depth_to_normals.py new file mode 100644 index 000000000..8bb28bd4f --- /dev/null +++ b/worldgen/tools/ground_truth/depth_to_normals.py @@ -0,0 +1,72 @@ +# Copyright (c) Princeton University. +# This source code is licensed under the BSD 3-Clause license found in the LICENSE file in the root directory of this source tree. + +# Authors: Lahav Lipson + +import argparse +import json +from pathlib import Path + +import numpy as np +from einops import einsum +from imageio.v3 import imread, imwrite +from numpy.linalg import inv + + +def transform(T, p): + assert T.shape == (4,4) + p = T[:3,:3] @ p + return p + T[:3, [3]] + +def from_homog(x): + return x[:-1] / x[[-1]] + +def unproject(depth, K): + H, W = depth.shape + x, y = np.meshgrid(np.arange(W), np.arange(H), indexing='xy') + img_coords = np.stack((x, y, np.ones_like(x)), axis=-1).astype(np.float64) + return einsum(depth, img_coords, inv(K), 'H W, H W j, i j -> H W i') + +def normalize(v): + return v / np.linalg.norm(v, axis=-1, keepdims=True) + +if __name__ == "__main__": + + parser = argparse.ArgumentParser() + parser.add_argument('folder', type=Path) + parser.add_argument('frame', type=int) + parser.add_argument('--output', type=Path, default=Path("testbed")) + args = parser.parse_args() + + folder_data = json.loads((args.folder / "summary.json").read_text()) + + depth_paths = folder_data["Depth"]['npy']["00"]["00"] + image_paths = folder_data["Image"]['png']["00"]["00"] + Ks = folder_data["Camera Intrinsics"]['npy']["00"]["00"] + + frame = f"{args.frame:04d}" + + image = imread(args.folder / image_paths[frame]) + depth = np.load(args.folder / depth_paths[frame]) + print(depth_paths) + K1 = np.load(args.folder / Ks[frame]) + cam_coords = unproject(depth, K1) + + cam_coords = cam_coords * np.array([1., -1., -1]) + + mask = ~np.isinf(depth) + depth[~mask] = -1 + + vy = normalize(cam_coords[1:,1:] - cam_coords[:-1,1:]) + vx = normalize(cam_coords[1:,1:] - cam_coords[1:,:-1]) + cross_prod = np.cross(vy, vx) + normals = normalize(cross_prod) + print(cross_prod.shape, mask.shape) + normals[~mask[1:,1:]] = 0 + + normals_color = np.round((normals + 1) * (255/2)).astype(np.uint8) + + imwrite(args.output / "A.png", image) + print(f'Wrote {args.output / "A.png"}') + imwrite(args.output / "B.png", normals_color) + print(f'Wrote {args.output / "B.png"}') \ No newline at end of file diff --git a/worldgen/tools/ground_truth/optical_flow_warp.py b/worldgen/tools/ground_truth/optical_flow_warp.py index 724379f2d..a21d3e686 100644 --- a/worldgen/tools/ground_truth/optical_flow_warp.py +++ b/worldgen/tools/ground_truth/optical_flow_warp.py @@ -13,9 +13,6 @@ if __name__ == "__main__": - import os,sys - sys.path.append(os.getcwd()) - parser = argparse.ArgumentParser() parser.add_argument('folder', type=Path) parser.add_argument('frame', type=int) diff --git a/worldgen/tools/ground_truth/rigid_warp.py b/worldgen/tools/ground_truth/rigid_warp.py index 7ac191440..18d4a007f 100644 --- a/worldgen/tools/ground_truth/rigid_warp.py +++ b/worldgen/tools/ground_truth/rigid_warp.py @@ -9,33 +9,30 @@ import cv2 import numpy as np +from einops import einsum from imageio.v3 import imread, imwrite from numpy.linalg import inv + def transform(T, p): assert T.shape == (4,4) - p = T[:3,:3] @ p - return p + T[:3, [3]] + return einsum(p, T[:3,:3], 'H W j, i j -> H W i') + T[:3, 3] def from_homog(x): - return x[:-1] / x[[-1]] + return x[...,:-1] / x[...,[-1]] def reproject(depth1, pose1, pose2, K1, K2): H, W = depth1.shape x, y = np.meshgrid(np.arange(W), np.arange(H), indexing='xy') - img_1_coords = np.stack((x, y, np.ones_like(x)), axis=-1).reshape((H*W, 3)).T.astype(np.float64) + img_1_coords = np.stack((x, y, np.ones_like(x)), axis=-1).astype(np.float64) + cam1_coords = einsum(depth1, img_1_coords, inv(K1), 'H W, H W j, i j -> H W i') rel_pose = inv(pose2) @ pose1 - cam1_coords = depth1.reshape((1, H*W)) * (inv(K1) @ img_1_coords) cam2_coords = transform(rel_pose, cam1_coords) - img2_coords = from_homog(K2 @ cam2_coords) - return img2_coords.T.reshape((H, W, 2)) + return from_homog(einsum(cam2_coords, K2, 'H W j, i j -> H W i')) if __name__ == "__main__": - import os,sys - sys.path.append(os.getcwd()) - parser = argparse.ArgumentParser() parser.add_argument('folder', type=Path) parser.add_argument('frame_1', type=int) diff --git a/worldgen/tools/ground_truth/segmentation_lookup.py b/worldgen/tools/ground_truth/segmentation_lookup.py index cd7340221..b8f8e809a 100644 --- a/worldgen/tools/ground_truth/segmentation_lookup.py +++ b/worldgen/tools/ground_truth/segmentation_lookup.py @@ -49,6 +49,7 @@ def compute_boxes(indices, binary_tag_mask): if __name__ == "__main__": + parser = argparse.ArgumentParser() parser.add_argument('folder', type=Path) parser.add_argument('frame', type=int) diff --git a/worldgen/tools/summarize.py b/worldgen/tools/summarize.py index b39f8fda3..56bf24193 100644 --- a/worldgen/tools/summarize.py +++ b/worldgen/tools/summarize.py @@ -39,7 +39,7 @@ def summarize_folder(base_folder): output = defaultdict(make_defaultdict(make_defaultdict(make_defaultdict(dict)))) max_frame = -1 for file_path in base_folder.rglob('*'): - if not file_path.is_file: + if (not file_path.is_file) or ("saved_mesh" in file_path.parts): continue if match := re.fullmatch("(.*)_([0-9]{4})_([0-9]{2})_([0-9]{2})\.([a-z]+)", file_path.name): diff --git a/worldgen/util/exporting.py b/worldgen/util/exporting.py index e5663c642..f66344228 100644 --- a/worldgen/util/exporting.py +++ b/worldgen/util/exporting.py @@ -171,9 +171,6 @@ def save_obj_and_instances(output_folder, previous_frame_mesh_id_mapping, curren if object_name not in object_names_mapping: object_names_mapping[object_name] = len(object_names_mapping) + 1 - group_name = parse_group_from_name(object_name) - semantic = parse_semantic_from_name(object_name) - # Flush the .npz to avoid OOM if (len(npz_data) > 0) and ((running_total_verts + current_obj_num_verts) >= MAX_NUM_VERTS): np.savez(filename, **npz_data) @@ -212,7 +209,12 @@ def save_obj_and_instances(output_folder, previous_frame_mesh_id_mapping, curren npz_data[f"{mesh_id}_instance_ids"] = instance_ids_array obj = bpy.data.objects[object_name] json_val = {"filename": filename.name, "mesh_id": mesh_id, "object_name": object_name, "num_verts": current_obj_num_verts, "children": [], - "object_type": obj.type, "semantic": semantic, "group": group_name, "num_instances": matrices.shape[0], "object_idx": object_names_mapping[object_name]} + "object_type": obj.type, "num_instances": matrices.shape[0], "object_idx": object_names_mapping[object_name]} + if obj.type == "MESH": + json_val['num_verts'] = len(obj.data.vertices) + json_val['num_faces'] = len(obj.data.polygons) + json_val['materials'] = obj.material_slots.keys() + json_val['unapplied_modifiers'] = obj.modifiers.keys() if not is_instance: non_aa_bbox = np.asarray([(obj.matrix_world @ mathutils.Vector(v)) for v in obj.bound_box], dtype=np.float32) json_val["instance_bbox"] = calc_aa_bbox(non_aa_bbox).tolist() @@ -220,8 +222,10 @@ def save_obj_and_instances(output_folder, previous_frame_mesh_id_mapping, curren else: combined_bbox, instance_bbox = calc_instance_bbox(matrices, item["vertex_lookup"]) json_val.update({"bbox": combined_bbox.tolist(), "instance_bbox": instance_bbox.tolist()}) - if obj.parent is not None: - json_val["parent"] = obj.parent.name + for child_obj in obj.children: + if child_obj.name not in object_names_mapping: + object_names_mapping[child_obj.name] = len(object_names_mapping) + 1 + json_val["children"].append(object_names_mapping[child_obj.name]) json_data.append(json_val) running_total_verts += current_obj_num_verts @@ -233,25 +237,16 @@ def save_obj_and_instances(output_folder, previous_frame_mesh_id_mapping, curren for obj in bpy.data.objects: if obj.type not in {"MESH", "CURVES", "CAMERA"}: object_name = obj.name - group_name = parse_group_from_name(object_name) - semantic = parse_semantic_from_name(object_name) if object_name not in object_names_mapping: object_names_mapping[object_name] = len(object_names_mapping) + 1 non_aa_bbox = np.asarray([(obj.matrix_world @ mathutils.Vector(v)) for v in obj.bound_box]) - json_val = {"object_name": object_name, "object_type": obj.type, "semantic": semantic, "children": [], - "group": group_name, "bbox": calc_aa_bbox(non_aa_bbox).tolist(), "object_idx": object_names_mapping[object_name]} - if obj.parent is not None: - json_val["parent"] = obj.parent.name + json_val = {"object_name": object_name, "object_type": obj.type, "children": [], + "bbox": calc_aa_bbox(non_aa_bbox).tolist(), "object_idx": object_names_mapping[object_name]} + for child_obj in obj.children: + if child_obj.name not in object_names_mapping: + object_names_mapping[child_obj.name] = len(object_names_mapping) + 1 + json_val["children"].append(object_names_mapping[child_obj.name]) json_data.append(json_val) - for obj in json_data: - if 'parent' in obj: - for p in json_data: - if p['object_name'] == obj['parent']: - p['children'] = list(set(p['children']).union({obj['object_name']})) - for key, val in list(obj.items()): - if val is None: - del obj[key] - # Save JSON (output_folder / "saved_mesh.json").write_text(json.dumps(json_data, indent=4))