Source code for blenderproc.python.utility.LabelIdMapping

""" Provides a class to handle the mapping between different semantic segmentations. """

import csv


[docs] class LabelIdMapping: """ Handles category id mapping for semantic segmentation maps. """ def __init__(self): # maps from an id to its name. E.g. id_label_map[0] = "void" self._id_label_map = {} # maps a class/category name to its id. E.g. label_id_map["void"] = 0 self._label_id_map = {} self._num_ids = 0
[docs] @staticmethod def from_csv(path, label_col_name="name", id_col_name="id") -> "LabelIdMapping": """ Builds a label-id mapping based on the given csv file. :param path: The path to a csv file. :param label_col_name: The name of the column which should be used as label. :param id_col_name: The name of the column which should be used as id. :return: The built label mapping object. """ with open(path, 'r', encoding="utf-8") as csv_file: reader = csv.DictReader(csv_file) mapping = LabelIdMapping() for row in reader: mapping.add(row[label_col_name], int(row[id_col_name])) return mapping
[docs] @staticmethod def from_dict(label_to_id: dict) -> "LabelIdMapping": """ Builds a label-id mapping based on the given dict. :param label_to_id: A dict where keys are labels and values are ids. :return: The built label mapping object. """ mapping = LabelIdMapping() for label, id_value in label_to_id.items(): mapping.add(label, id_value) return mapping
[docs] def add(self, label: str, id_value: int): """ Inserts the given label-id pair into the mapping. :param label: The label of the pair. :param id_value: The id of the pair """ if self.has_id(id_value): raise Exception("There already exists a label-id mapping for the id " + str(id_value)) if self.has_label(label): raise Exception("There already exists a label-id mapping for the label " + label) self._id_label_map[id_value] = label self._label_id_map[label] = id_value self._num_ids = max(self._num_ids, id_value + 1)
[docs] def id_from_label(self, label: str) -> int: """ Returns the id assigned to the given label. :param label: The label to look for. :return: The id with the given label. """ return self._label_id_map[label]
[docs] def label_from_id(self, id_value: int) -> str: """ Returns the label assigned to the given id. :param id_value: The id to look for. :return: The label with the given id. """ return self._id_label_map[id_value]
[docs] def has_label(self, label: str) -> bool: """ Checks if the mapping contains the given label. :param label: The label to look for. :return: True, if the label is already in use. """ return label in self._label_id_map
[docs] def has_id(self, id_value: int) -> bool: """ Checks if the mapping contains the given id. :param id_value: The id to look for. :return: True, if the id is already in use. """ return id_value in self._id_label_map