diff --git a/omnigibson/sensors/vision_sensor.py b/omnigibson/sensors/vision_sensor.py index b53cdb55c..7ee678369 100644 --- a/omnigibson/sensors/vision_sensor.py +++ b/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 diff --git a/omnigibson/utils/vision_utils.py b/omnigibson/utils/vision_utils.py index 219c7819e..d23bdebab 100644 --- a/omnigibson/utils/vision_utils.py +++ b/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]