Only set colors for enabled objects (#4936)

* Only create colormap for enabled labels

* Fix assigning
This commit is contained in:
Nicolas Mowen
2023-01-06 18:31:54 -07:00
committed by GitHub
parent d9c45a76fe
commit 543cad5497
2 changed files with 14 additions and 4 deletions

View File

@@ -55,11 +55,13 @@ class ModelConfig(BaseModel):
**load_labels(config.get("labelmap_path", "/labelmap.txt")),
**config.get("labelmap", {}),
}
cmap = plt.cm.get_cmap("tab10", len(self._merged_labelmap.keys()))
self._colormap = {}
for key, val in self._merged_labelmap.items():
def create_colormap(self, enabled_labels: set[str]) -> None:
"""Get a list of colors for enabled labels."""
cmap = plt.cm.get_cmap("tab10", len(enabled_labels))
for key, val in enumerate(enabled_labels):
self._colormap[val] = tuple(int(round(255 * c)) for c in cmap(key)[:3])
class Config: