How to write a custom writer to save depth image?

i try to modify the MyCustomWriter in tutorial to write depth information to .npy file

import numpy as np
from omni.replicator.core import AnnotatorRegistry, BackendDispatch, Writer, WriterRegistry

class MyCustomWriter(Writer):
    def __init__(
        self,
        output_dir,
        rgb = True,
        distance_to_camera = True,
        normals = False,
    ):
        self.version = "0.0.1"
        self.backend = BackendDispatch({"paths": {"out_dir": output_dir}})
        if rgb:
            self.annotators.append(AnnotatorRegistry.get_annotator("rgb"))
        if normals:
            self.annotators.append(AnnotatorRegistry.get_annotator("normals"))
        if distance_to_camera:
            self.annotators.append(AnnotatorRegistry.get_annotator("distance_to_camera"))
        self._frame_id = 0
        self.max_xl = 0
        self.min_xl = 0

    def write(self, data: dict):
        for annotator in data.keys():
            # If there are multiple render products the data will be stored in subfolders
            annotator_split = annotator.split("-")
            render_product_path = ""
            multi_render_prod = 0
            if len(annotator_split) > 1:
                multi_render_prod = 1
                render_product_name = annotator_split[-1]
                render_product_path = f"{render_product_name}/"

            # rgb
            if annotator.startswith("rgb"):
                if multi_render_prod:
                    render_product_path += "rgb/"
                filename = f"{render_product_path}rgb_{self._frame_id}.png"
                if((self.min_xl<-0.1) or (self.max_xl>0.1) or (self.max_xl-self.min_xl>0.15)):
                    filename = "abnormal_" + filename
                else:
                    filename = "normal_" + filename
                print(f"[{self._frame_id}] Writing {self.backend.output_dir}/{filename} ..")
                self.backend.write_image(filename, data[annotator])

            # semantic_segmentation
            if annotator.startswith("normals"):
                if multi_render_prod:
                    render_product_path += "normals/"
                filename = f"{render_product_path}normals_{self._frame_id}.png"
                print(f"[{self._frame_id}] Writing {self.backend.output_dir}/{filename} ..")
                colored_data = ((data[annotator] * 0.5 + 0.5) * 255).astype(np.uint8)
                self.backend.write_image(filename, colored_data)

            # distance
            if annotator.startswith("distance"):
                if multi_render_prod:
                    render_product_path += "distances/"
                filename = f"{render_product_path}distance_{self._frame_id}.npy"
                print(f"[{self._frame_id}] Writing {self.backend.output_dir}/{filename} ..")
                self.backend.write_depth(filename, data[annotator])
        self._frame_id += 1

    def on_final_frame(self):
        self._frame_id = 0
    
    def set_max_min_xl(self, max_xl, min_xl):
        self.max_xl = max_xl
        self.min_xl = min_xl

WriterRegistry.register(MyCustomWriter)

but i can’t use write_image to save depth information, so i tried to use write_depth, but it say BackendDispatch has no attribute ‘wirte_depth’.

2024-03-14 03:02:06 [3,386,607ms] [Error] [omni.graph.core.plugin] /Render/PostProcess/SDGPipeline/RenderProduct_View_MyCustomWriterWriter: [/Render/PostProcess/SDGPipeline] Assertion raised in compute - 'BackendDispatch' object has no attribute 'write_depth'
  File "/isaac-sim/extscache/omni.replicator.core-1.10.20+105.1.lx64.r.cp310/omni/replicator/core/ogn/python/_impl/nodes/OgnWriter.py", line 116, in compute
    writer._write(writer_payload)
  File "/isaac-sim/extscache/omni.replicator.core-1.10.20+105.1.lx64.r.cp310/omni/replicator/core/scripts/writers.py", line 325, in _write
    self.write(data)
  File "/tmp/carb.Zc1tcB/script_1710385302.py", line 62, in write
    self.backend.write_depth(filename, data[annotator])

which function should i call so i can save depth information in .npy file?

Hi there,

here are the functions from the basicwriter (/omni/replicator/core/scripts/writers_default/basicwriter.py) writing the depth data as npy format:

if annotator.startswith("distance_to_camera"):
    if multi_render_prod:
        render_product_path += "distance_to_camera/"
    self._write_distance_to_camera(data, render_product_path, annotator)
  
if annotator.startswith("distance_to_image_plane"):
    if multi_render_prod:
        render_product_path += "distance_to_image_plane/"
    self._write_distance_to_image_plane(data, render_product_path, annotator)

[..]

def _write_distance_to_camera(self, data: dict, render_product_path: str, annotator: str):
    dist_to_cam_data = data[annotator]
    file_path = (
        f"{render_product_path}distance_to_camera_{self._sequence_id}{self._frame_id:0{self._frame_padding}}.npy"
    )
    self._backend.schedule(F.write_np, data=dist_to_cam_data, path=file_path)

def _write_distance_to_image_plane(self, data: dict, render_product_path: str, annotator: str):
    dis_to_img_plane_data = data[annotator]
    file_path = f"{render_product_path}distance_to_image_plane_{self._sequence_id}{self._frame_id:0{self._frame_padding}}.npy"
    self._backend.schedule(F.write_np, data=dis_to_img_plane_data, path=file_path)

you can however change the functions and write the data in different formats as well.