120 lines
4.6 KiB
Python
120 lines
4.6 KiB
Python
import numpy as np
|
|
import omni.replicator.core as rep
|
|
from omni.isaac.core.prims import BaseSensor
|
|
from omni.isaac.core.prims.xform_prim import XFormPrim
|
|
from pxr import UsdGeom
|
|
from scipy.spatial.transform import Rotation as R
|
|
|
|
|
|
class AnnotatorManager:
|
|
def __init__(self, world):
|
|
self.annotators = {}
|
|
self.render_products = {}
|
|
self.cameras = {}
|
|
self.annotator_types = [
|
|
"rgb",
|
|
"distance_to_camera",
|
|
"pointcloud",
|
|
"semantic_segmentation",
|
|
]
|
|
self.camera_prims = {}
|
|
self.camera_info = {}
|
|
self.world = world
|
|
|
|
def registerCamera(
|
|
self, parent_prim, name, translation, orientation, resolution=(1024, 600)
|
|
):
|
|
"""
|
|
Register a camera in the scene. The camera will be attached to the parent_prim.
|
|
parent_prim: The prim path of the parent prim. The camera will be attached to this prim.
|
|
name: The name of the camera.
|
|
translation: The translation of the camera relative to the parent prim in the form of (x, y, z).
|
|
orientation: The orientation of the camera relative to the parent prim in the form of (x, y, z, w).
|
|
resolution: The resolution of the camera in the form of (width, height).
|
|
"""
|
|
qx, qy, qz, qw = tuple(orientation)
|
|
xformprim = XFormPrim(
|
|
prim_path=parent_prim + "/" + name,
|
|
name=name,
|
|
translation=tuple(translation),
|
|
orientation=(qw, qx, qy, qz), # Omniverse core convention (w, x, y, z)
|
|
)
|
|
self.camera_prims[name] = xformprim
|
|
self.cameras[name] = self.world.stage.DefinePrim(
|
|
xformprim.prim_path + "/" + name, "Camera"
|
|
)
|
|
UsdGeom.Xformable(self.cameras[name]).AddTranslateOp().Set((0.0, 0.0, 0.0))
|
|
# Rotate around x for 180 degrees to make the convention consistent with opencv/ROS
|
|
UsdGeom.Xformable(self.cameras[name]).AddRotateXYZOp().Set((180.0, 0.0, 0.0))
|
|
self.camera_info[name] = {
|
|
"translation": translation,
|
|
"orientation": orientation,
|
|
"resolution": resolution,
|
|
}
|
|
self.render_products[name] = rep.create.render_product(
|
|
str(self.cameras[name].GetPrimPath()), resolution
|
|
)
|
|
|
|
def registerAnnotator(self, type, camera_name):
|
|
assert (
|
|
type in self.annotator_types
|
|
), "Requested replicator type not available. Choose from: " + str(
|
|
self.annotator_types
|
|
)
|
|
assert (
|
|
camera_name in self.cameras.keys()
|
|
), "The requested camera is not registered. Please register the camera first."
|
|
self.annotators[camera_name + ":" + type] = rep.AnnotatorRegistry.get_annotator(
|
|
type
|
|
)
|
|
self.annotators[camera_name + ":" + type].attach(
|
|
[self.render_products[camera_name]]
|
|
)
|
|
|
|
def getData(self, annotator_name):
|
|
assert (
|
|
annotator_name in self.annotators.keys()
|
|
), "The requested annotator is not registered. Please register the annotator first."
|
|
return self.annotators[annotator_name].get_data()
|
|
|
|
def getStreams(self):
|
|
return list(self.annotators.keys())
|
|
|
|
def setFocalLength(self, name, value):
|
|
assert (
|
|
name in self.cameras.keys()
|
|
), "The requested camera is not registered. Please register the camera first."
|
|
self.cameras[name].GetAttribute("focalLength").Set(value)
|
|
|
|
def setClippingRange(self, name, min=0.2, max=1000000.0):
|
|
assert (
|
|
name in self.cameras.keys()
|
|
), "The requested camera is not registered. Please register the camera first."
|
|
self.cameras[name].GetAttribute("clippingRange").Set((min, max))
|
|
|
|
# TODO: implement these functions
|
|
def getCameraIntrinsics(self, name):
|
|
return None
|
|
|
|
def getCameraExtrinsics(self, name):
|
|
return None
|
|
|
|
def getCameraPose(self, name):
|
|
"""
|
|
Get the camera pose in the world frame
|
|
@param name: (str) The name of the camera
|
|
@return: (t, q) The camera pose (np.array([x, y, z]), np.array([qw, qx, qy, qz]))
|
|
"""
|
|
t, q = self.camera_prims[name].get_world_pose()
|
|
return t, q
|
|
|
|
def setCameraPose(self, name, t, q):
|
|
"""
|
|
Get the camera pose in the world frame
|
|
@param name: (str) The name of the camera
|
|
@param t: (mp.array([x, y, z])) The translation of the camera in the world frame
|
|
@param q: (mp.array([qw, qx, qy, qz])) The orientation of the camera in the world frame
|
|
@return: None
|
|
"""
|
|
self.camera_prims[name].set_local_pose(t, q)
|