Merge pull request #667 from StanfordVL/replicator-hot-fix

Replicator hot fix

omnigibson/sensors/vision_sensor.py

@@ -382,8 +382,16 @@ def _remap_instance_segmentation(self, img, id_to_labels, semantic_img, semantic
                 # If the category name is not in the registered systems, 
                 # which happens because replicator sometimes returns segmentation map and id_to_labels that are not in sync,
                 # we will label this as "unlabelled" for now
+                # This only happens with a very small number of pixels, e.g. 0.1% of the image
                 else:
+                    num_of_pixels = len(np.where(img == key)[0])
+                    resolution = (self._load_config["image_width"], self._load_config["image_height"])
+                    percentage = (num_of_pixels / (resolution[0] * resolution[1])) * 100
+                    if percentage > 2:
+                        og.log.warning(f"Marking {category_name} as unlabelled due to image & id_to_labels mismatch!"
+                                       f"Percentage of pixels: {percentage}%")
                     value = "unlabelled"
+                    self._register_instance(value, id=id)
                 replicator_mapping[key] = value
 
         registry = VisionSensor.INSTANCE_ID_REGISTRY if id else VisionSensor.INSTANCE_REGISTRY

omnigibson/utils/vision_utils.py

@@ -2,6 +2,8 @@
 
 import numpy as np
 from PIL import Image, ImageDraw
+import omnigibson as og
+from omnigibson.utils.constants import semantic_class_name_to_id
 
 try:
     import accimage
@@ -66,6 +68,7 @@ class Remapper:
     def __init__(self):
         self.key_array = np.array([], dtype=np.uint32)  # Initialize the key_array as empty
         self.known_ids = set()
+        self.warning_printed = set()
 
     def clear(self):
         """Resets the key_array to empty."""
@@ -137,7 +140,11 @@ def remap_bbox(self, semantic_id):
         Returns:
             int: The remapped id.
         """
-        assert semantic_id < len(self.key_array), f"Semantic id {semantic_id} is out of range!"
+        if semantic_id >= len(self.key_array):
+            if semantic_id not in self.warning_printed:
+                og.log.warning(f"We do not have semantic information about bounding box semantic id {semantic_id} yet. Marking as unlabelled.")
+                self.warning_printed.add(semantic_id)
+            return semantic_class_name_to_id()['unlabelled']
         return self.key_array[semantic_id]