From 421462fb2ea0a1f1eac0480efd35865b30a9a7c4 Mon Sep 17 00:00:00 2001
From: hang-yin <yinhang0226@gmail.com>
Date: Thu, 3 Oct 2024 15:41:52 -0700
Subject: [PATCH] Fix particle segmentation with viewer camera & no particle
 systems in instance segmentation

---
 omnigibson/sensors/vision_sensor.py | 49 ++++++++++++++---------------
 omnigibson/systems/system_base.py   |  2 ++
 2 files changed, 26 insertions(+), 25 deletions(-)

diff --git a/omnigibson/sensors/vision_sensor.py b/omnigibson/sensors/vision_sensor.py
index bba636a19..17b68b029 100644
--- a/omnigibson/sensors/vision_sensor.py
+++ b/omnigibson/sensors/vision_sensor.py
@@ -7,6 +7,7 @@ import torch as th
 import omnigibson as og
 import omnigibson.lazy as lazy
 from omnigibson.sensors.sensor_base import BaseSensor
+from omnigibson.systems.system_base import get_all_system_names
 from omnigibson.utils.constants import (
     MAX_CLASS_COUNT,
     MAX_INSTANCE_COUNT,
@@ -353,9 +354,7 @@ class VisionSensor(BaseSensor):
             if "," in replicator_mapping[key]:
                 # If there are multiple class names, grab the one that is a registered system
                 # This happens with MacroVisual particles, e.g. {"11": {"class": "breakfast_table,stain"}}
-                categories = [
-                    cat for cat in replicator_mapping[key].split(",") if cat in self.scene.available_systems.keys()
-                ]
+                categories = [cat for cat in replicator_mapping[key].split(",") if cat in get_all_system_names()]
                 assert (
                     len(categories) == 1
                 ), "There should be exactly one category that belongs to scene.system_registry"
@@ -415,30 +414,30 @@ class VisionSensor(BaseSensor):
             if value in ["BACKGROUND", "UNLABELLED"]:
                 value = value.lower()
             elif "/" in value:
-                prim_name = value.split("/")[-1]
-                # Hacky way to get the particles of MacroVisual/PhysicalParticleSystem
-                # Remap instance segmentation and instance segmentation ID labels to system name
-                if "Particle" in prim_name:
-                    category_name = prim_name.split("Particle")[0]
-                    assert (
-                        category_name in self.scene.available_systems.keys()
-                    ), f"System name {category_name} is not in the registered systems!"
-                    value = category_name
-                else:
+                if not id:
                     # Remap instance segmentation labels to object name
-                    if not id:
-                        # value is the prim path of the object
-                        if og.sim.floor_plane is not None and value == og.sim.floor_plane.prim_path:
-                            value = "groundPlane"
-                        else:
-                            obj = self.scene.object_registry("prim_path", value)
-                            # Remap instance segmentation labels from prim path to object name
-                            assert obj is not None, f"Object with prim path {value} cannot be found in objct registry!"
-                            value = obj.name
-
-                    # Keep the instance segmentation ID labels intact (prim paths of visual meshes)
+                    if og.sim.floor_plane is not None and value == og.sim.floor_plane.prim_path:
+                        value = "groundPlane"
                     else:
-                        pass
+                        obj = None
+                        if self.scene is not None:
+                            # If this is a camera within a scene, we check the object registry of the scene
+                            obj = self.scene.object_registry("prim_path", value)
+                        else:
+                            # If this is the viewer camera, we check each object registry
+                            for scene in og.sim.scenes:
+                                obj = scene.object_registry("prim_path", value)
+                                if obj:
+                                    break
+                        if obj is not None:
+                            # This is an object, so we remap the instance segmentation label to the object name
+                            value = obj.name
+                        else:
+                            # This is a particle system, we remap them to unlabelled (will fix this in a future release)
+                            value = "unlabelled"
+                # Keep the instance segmentation ID labels intact (prim paths of visual meshes)
+                else:
+                    pass
             else:
                 # TODO: This is a temporary fix unexpected labels e.g. INVALID introduced in new Isaac Sim versions
                 value = "unlabelled"
diff --git a/omnigibson/systems/system_base.py b/omnigibson/systems/system_base.py
index 43e5dea9a..3e45e061a 100644
--- a/omnigibson/systems/system_base.py
+++ b/omnigibson/systems/system_base.py
@@ -1,5 +1,6 @@
 import json
 import os
+from functools import cache
 
 import torch as th
 
@@ -957,6 +958,7 @@ class PhysicalParticleSystem(BaseSystem):
         return success
 
 
+@cache
 def get_all_system_names():
     """
     Gets all available systems from the OmniGibson dataset