Base class for standard datasets. Args: roots (str): paths to the dataset
| 10 | |
| 11 | |
| 12 | class StandardDatasetBase(Dataset): |
| 13 | """ |
| 14 | Base class for standard datasets. |
| 15 | |
| 16 | Args: |
| 17 | roots (str): paths to the dataset |
| 18 | """ |
| 19 | |
| 20 | def __init__(self, |
| 21 | roots: str, |
| 22 | ): |
| 23 | super().__init__() |
| 24 | self.roots = roots.split(',') |
| 25 | self.instances = [] |
| 26 | self.metadata = pd.DataFrame() |
| 27 | |
| 28 | self._stats = {} |
| 29 | for root in self.roots: |
| 30 | key = os.path.basename(root) |
| 31 | self._stats[key] = {} |
| 32 | metadata = pd.read_csv(os.path.join(root, 'metadata.csv')) |
| 33 | self._stats[key]['Total'] = len(metadata) |
| 34 | metadata, stats = self.filter_metadata(metadata) |
| 35 | self._stats[key].update(stats) |
| 36 | self.instances.extend([(root, sha256) for sha256 in metadata['sha256'].values]) |
| 37 | metadata.set_index('sha256', inplace=True) |
| 38 | self.metadata = pd.concat([self.metadata, metadata]) |
| 39 | |
| 40 | @abstractmethod |
| 41 | def filter_metadata(self, metadata: pd.DataFrame) -> Tuple[pd.DataFrame, Dict[str, int]]: |
| 42 | pass |
| 43 | |
| 44 | @abstractmethod |
| 45 | def get_instance(self, root: str, instance: str) -> Dict[str, Any]: |
| 46 | pass |
| 47 | |
| 48 | def __len__(self): |
| 49 | return len(self.instances) |
| 50 | |
| 51 | def __getitem__(self, index) -> Dict[str, Any]: |
| 52 | try: |
| 53 | root, instance = self.instances[index] |
| 54 | return self.get_instance(root, instance) |
| 55 | except Exception as e: |
| 56 | print(e) |
| 57 | return self.__getitem__(np.random.randint(0, len(self))) |
| 58 | |
| 59 | def __str__(self): |
| 60 | lines = [] |
| 61 | lines.append(self.__class__.__name__) |
| 62 | lines.append(f' - Total instances: {len(self)}') |
| 63 | lines.append(f' - Sources:') |
| 64 | for key, stats in self._stats.items(): |
| 65 | lines.append(f' - {key}:') |
| 66 | for k, v in stats.items(): |
| 67 | lines.append(f' - {k}: {v}') |
| 68 | return '\n'.join(lines) |
| 69 |
nothing calls this directly
no outgoing calls
no test coverage detected