INNER CODE UNIT · Python

__getitem__

allenai/ai2thor · ai2thor/server.py:311

    def __getitem__(self, cls: str) -> Optional[Tuple[Tuple[int, int, int, int], ...]]:
        if cls in self._detections2d:
            return self._detections2d[cls]

        detections = []

        for color in self.instance_masks.class_colors.get(cls, []):
            mask = (
                self.instance_masks.instance_segmentation_frame_uint32
                == self.instance_masks._integer_color_key(color)
            )
            bb = self.mask_bounding_box(mask)
            if bb:
                detections.append(bb)
        if detections:
            self._detections2d[cls] = tuple(detections)
        else:
            raise KeyError(cls)

View source record →