Spaces:
Running on Zero
Running on Zero
| from abc import ABC, abstractmethod | |
| from typing import Dict, List, Optional, Union | |
| class DatasetInfo(ABC): | |
| def __init__(self): | |
| pass | |
| def num_classes(self): | |
| pass | |
| def label_names(self): | |
| pass | |
| def label_descriptions(self, detailed: bool = False, as_dict: bool = False) -> Union[List[str], Dict[str, str]]: | |
| pass | |
| def index_to_label_name(self, index) -> str: | |
| pass | |
| def index_to_description(self, index: int, detailed: bool = False) -> str: | |
| pass | |
| def label_name_to_description(self, label: str, detailed: bool = False) -> str: | |
| pass | |
| class CustomDatasetInfo(DatasetInfo): | |
| """ DatasetInfo that wraps passed values for custom datasets.""" | |
| def __init__( | |
| self, | |
| label_names: Union[List[str], Dict[int, str]], | |
| label_descriptions: Optional[Dict[str, str]] = None | |
| ): | |
| super().__init__() | |
| assert len(label_names) > 0 | |
| self._label_names = label_names # label index => label name mapping | |
| self._label_descriptions = label_descriptions # label name => label description mapping | |
| if self._label_descriptions is not None: | |
| # validate descriptions (label names required) | |
| assert isinstance(self._label_descriptions, dict) | |
| for n in self._label_names: | |
| assert n in self._label_descriptions | |
| def num_classes(self): | |
| return len(self._label_names) | |
| def label_names(self): | |
| return self._label_names | |
| def label_descriptions(self, detailed: bool = False, as_dict: bool = False) -> Union[List[str], Dict[str, str]]: | |
| return self._label_descriptions | |
| def label_name_to_description(self, label: str, detailed: bool = False) -> str: | |
| if self._label_descriptions: | |
| return self._label_descriptions[label] | |
| return label # return label name itself if a descriptions is not present | |
| def index_to_label_name(self, index) -> str: | |
| assert 0 <= index < len(self._label_names) | |
| return self._label_names[index] | |
| def index_to_description(self, index: int, detailed: bool = False) -> str: | |
| label = self.index_to_label_name(index) | |
| return self.label_name_to_description(label, detailed=detailed) | |