Skip to content

Commit

Permalink
Adding render modalities
Browse files Browse the repository at this point in the history
  • Loading branch information
@abhihjoshi
abhihjoshi committed Dec 18, 2024
1 parent b59fa7a commit fb81686
Show file tree
Hide file tree
Showing 3 changed files with 97 additions and 16 deletions.
102 changes: 89 additions & 13 deletions robosuite/scripts/render_dataset_with_omniverse.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -74,6 +74,23 @@
"--keep_models", type=str, nargs="+", default=[], help="(optional) keep the model from the Mujoco XML file"
)

# adding rendering types
parser.add_argument(
"--rgb",
action="store_true",
default=False,
)
parser.add_argument(
"--normals",
action="store_true",
default=False,
)
parser.add_argument(
"--semantic_segmentation",
action="store_true",
default=False,
)

# Add arguments for launch
AppLauncher.add_app_launcher_args(parser)
# Parse the arguments
Expand All @@ -95,11 +112,13 @@
import cv2
import h5py
import lxml.etree as ET
import numpy as np
import omni
import omni.isaac.core.utils.stage as stage_utils
import omni.kit.app
import omni.replicator.core as rep
import omni.timeline
from pxr import Semantics
from termcolor import colored
from tqdm import tqdm

Expand All @@ -124,6 +143,8 @@
scene_option = mujoco.MjvOption()
scene_option.geomgroup = [0, 1, 0, 0, 0, 0]

render_modalities = {"rgb": args.rgb, "normals": args.normals, "semantic_segmentation": args.semantic_segmentation}


def make_sites_invisible(mujoco_xml):
"""
Expand Down Expand Up @@ -343,6 +364,24 @@ def link_env_with_ov(self):
)
exp.update_scene(data=data, scene_option=scene_option)
exp.add_light(pos=[0, 0, 0], intensity=1500, light_type="dome", light_name="dome_1")

# adds semantic information to objects in the scene
if args.semantic_segmentation:
for geom in exp.scene.geoms:
geom_id = geom.objid
geom_name = exp._get_geom_name(geom)
if geom_id in self.env.model._geom_ids_to_classes:
semantic_value = self.env.model._geom_ids_to_classes[geom_id]
if "site" in geom_name or "None" in geom_name:
continue
prim = exp.geom_refs[geom_name].usd_prim
instance_name = f"class_{semantic_value}"
sem = Semantics.SemanticsAPI.Apply(prim, instance_name)
sem.CreateSemanticTypeAttr()
sem.CreateSemanticDataAttr()
sem.GetSemanticTypeAttr().Set("class")
sem.GetSemanticDataAttr().Set(semantic_value)

return exp

def update_simulation(self, index):
Expand All @@ -369,6 +408,8 @@ def __init__(
output_dir: str = None,
image_output_format: str = "png",
rgb: bool = False,
normals: bool = False,
semantic_segmentation: bool = False,
frame_padding: int = 4,
):
self._output_dir = output_dir
Expand All @@ -385,9 +426,16 @@ def __init__(
self.data_structure = "annotator"
self.write_ready = False

# RGB
self.rgb = rgb
self.normals = normals
self.semantic_segmentation = semantic_segmentation

if rgb:
self.annotators.append(rep.AnnotatorRegistry.get_annotator("rgb"))
if normals:
self.annotators.append(rep.AnnotatorRegistry.get_annotator("normals"))
if semantic_segmentation:
self.annotators.append(rep.AnnotatorRegistry.get_annotator("semantic_segmentation", {"colorize": True}))

def write(self, data: dict):
"""Write function called from the OgnWriter node on every frame to process annotator output.
Expand All @@ -399,7 +447,25 @@ def write(self, data: dict):
for annotator_name, annotator_data in data["annotators"].items():
for idx, (render_product_name, anno_rp_data) in enumerate(annotator_data.items()):
if annotator_name == "rgb":
filepath = os.path.join(args.cameras[idx], f"rgb_{self._frame_id}.{self._image_output_format}")
filepath = os.path.join(
args.cameras[idx], "rgb", f"rgb_{self._frame_id}.{self._image_output_format}"
)
self._backend.write_image(filepath, anno_rp_data["data"])
elif annotator_name == "normals":
normals = anno_rp_data["data"][..., :3]
norm_lengths = np.linalg.norm(normals, axis=-1, keepdims=True)
normals_normalized = normals / np.clip(norm_lengths, 1e-8, None)
img = ((normals_normalized + 1) / 2 * 255).astype(np.uint8)
filepath = os.path.join(
args.cameras[idx], "normals", f"normals_{self._frame_id}.{self._image_output_format}"
)
self._backend.write_image(filepath, img)
elif annotator_name == "semantic_segmentation":
filepath = os.path.join(
args.cameras[idx],
"semantic_segmentation",
f"semantic_segmentation_{self._frame_id}.{self._image_output_format}",
)
self._backend.write_image(filepath, anno_rp_data["data"])

self._frame_id += 1
Expand Down Expand Up @@ -481,7 +547,12 @@ def init_recorder(self):

# Create writer for capturing generated data
self.writer = rep.WriterRegistry.get(self.writer_name)
self.writer.initialize(output_dir=self.output_dir, rgb=True)
self.writer.initialize(
output_dir=self.output_dir,
rgb=args.rgb,
normals=args.normals,
semantic_segmentation=args.semantic_segmentation,
)

print("Writer Initiazed")

Expand Down Expand Up @@ -589,22 +660,27 @@ def create_video_from_frames(self, frame_folder, output_path, fps=30):
video.release()
print(f"Video saved: {output_path}")

def create_video(self, videos_folder, camera, data_type):
camera_folder_path = os.path.join(self.output_dir, camera, data_type) # temp, change to render type
if not os.path.isdir(camera_folder_path):
return

# Construct output filename and path
output_filename = f"{camera}_{data_type}.mp4"
output_path = os.path.join(videos_folder, output_filename)

# Create the video from the frames in the camera folder
self.create_video_from_frames(camera_folder_path, output_path)

def process_folders(self):
videos_folder = os.path.join(self.output_dir, "videos")
os.makedirs(videos_folder, exist_ok=True)

# Iterate over each camera folder in the output directory
for camera in args.cameras:
camera_folder_path = os.path.join(self.output_dir, camera)
if not os.path.isdir(camera_folder_path):
continue

# Construct output filename and path
output_filename = f"{camera}_rgb.mp4"
output_path = os.path.join(videos_folder, output_filename)

# Create the video from the frames in the camera folder
self.create_video_from_frames(camera_folder_path, output_path)
for render_modality, selected in render_modalities.items():
if selected:
self.create_video(videos_folder=videos_folder, camera=camera, data_type=render_modality)


def main():
Expand Down
2 changes: 2 additions & 0 deletions robosuite/utils/usd/exporter.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -527,6 +527,8 @@ def _get_geom_name(self, geom) -> str:
geom_name = mujoco.mj_id2name(self.model, geom.objtype, geom.objid)
if not geom_name:
geom_name = "None"
geom_name = geom_name.replace("-", "m_")
geom_name = geom_name.replace("+", "p_")
geom_name += f"_id{geom.objid}"

# adding additional naming information to differentiate
Expand Down
9 changes: 6 additions & 3 deletions robosuite/utils/usd/objects.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -59,7 +59,10 @@ def __init__(
self.rgba = rgba
self.texture_file = texture_file

self.xform_path = f"/World/Mesh_Xform_{obj_name}"
self.obj_name = self.obj_name.replace("-", "m_")
self.obj_name = self.obj_name.replace("+", "p_")

self.xform_path = f"/World/Mesh_Xform_{self.obj_name}"
self.usd_xform = UsdGeom.Xform.Define(stage, self.xform_path)

# defining ops required by update function
Expand Down Expand Up @@ -199,7 +202,7 @@ def __init__(

self.dataid = dataid

mesh_path = f"{self.xform_path}/Mesh_{obj_name}"
mesh_path = f"{self.xform_path}/Mesh_{self.obj_name}"
self.usd_mesh = UsdGeom.Mesh.Define(stage, mesh_path)
self.usd_prim = stage.GetPrimAtPath(mesh_path)

Expand Down Expand Up @@ -288,7 +291,7 @@ def __init__(
self.mesh_config = mesh_config
self.prim_mesh = self.generate_primitive_mesh()

mesh_path = f"{self.xform_path}/Mesh_{obj_name}"
mesh_path = f"{self.xform_path}/Mesh_{self.obj_name}"
self.usd_mesh = UsdGeom.Mesh.Define(stage, mesh_path)
self.usd_prim = stage.GetPrimAtPath(mesh_path)

Expand Down

0 comments on commit fb81686

Please sign in to comment.