Hello,
I would like to create a Writer that for each image register save the position and rotation of the camera in a numpy array and save the numpy array at the end.
I code this :
class ArenaWriter(Writer):
def __init__(
self,
output_dir,
rig,
image_output_format="png",
n_data=10,
n_attr=5
):
self._output_dir = output_dir
self._backend = BackendDispatch({"paths": {"out_dir": output_dir}})
self._frame_id = 0
self._image_output_format = image_output_format
self._rig=rig
self._attributes = np.zeros(shape=(n_data, n_attr))
self.annotators = [AnnotatorRegistry.get_annotator("rgb")]
def write(self, data):
im_array = np.array(data['rgb'])
im_PIL = Image.fromarray(np.uint8(im_array))
im_PIL = im_PIL.resize((128,64), Image.ANTIALIAS)
im_array = np.array(im_PIL)
filepath = f"rgb_{self._frame_id}.{self._image_output_format}"
cam_pos = self._rig.get_world_pose()
pos = cam_pos[0]
rot = Rotation.from_quat(cam_pos[1])
rot_euler = rot.as_euler('xyz', degrees=True)
self._attributes[self._frame_id,:2] = pos[:2]
self._attributes[self._frame_id,2:] = cam_pos[1]
self._backend.write_image(filepath, im_array)
self._frame_id += 1
def on_final_frame(self):
np.save(self._output_dir + 'attributes.npy', self._attributes)
def main():
# Open the environment in a new stage
print(f"Loading Stage {ENV_URL}")
open_stage(ENV_URL)
stage = omni.usd.get_context().get_stage()
# Create Replicator Camera
cam = rep.create.camera(
position=(0, 0, 0.25),
rotation=(0, 0, 0),
focal_length=28,
fisheye_max_fov=110,
clipping_range=(0.01, 20.)
)
cam_node = cam.node
print(cam_node)
cam_rig_path = rep.utils.get_node_targets(cam_node, "inputs:prims")[0]
print(cam_rig_path)
cam_path = str(cam_rig_path) + "/Camera"
print(cam_path)
rig = XFormPrim(prim_path=cam_rig_path)
rep.WriterRegistry.register(ArenaWriter)
writer = rep.WriterRegistry.get("ArenaWriter")
out_dir = "/root/Documents/arena/data/"
writer.initialize(output_dir=out_dir, rig=rig, n_data=CONFIG["num_frames"])
# Create a Replicator render for the Isaac Sim API Camera
RESOLUTION = (1600, 1300)
camera_rp = rep.create.render_product(cam, RESOLUTION)
# Attach the render to the Writer
writer.attach([camera_rp])
with rep.trigger.on_frame():
with cam:
rep.modify.pose(
position=rep.distribution.uniform((-1.8, -1.3, 0.25), (1.8, 1.3, 0.25)),
rotation=rep.distribution.uniform((0, 0, 0), (0, 0, 360))
)
for i in range(CONFIG["num_frames"]):
rep.orchestrator.step()
writer.on_final_frame()
But I think this is suboptimal in term of time and I prefered to use :
rep.orchestrator.run()
# Wait until started
while not rep.orchestrator.get_is_started():
simulation_app.update()
# Wait until stopped
while rep.orchestrator.get_is_started():
simulation_app.update()
rep.BackendDispatch.wait_until_done()
rep.orchestrator.stop()
But with this I have some trouble to save camera position and rotation because there is a mismatch between camera parameters and the image saved.