A Loader for DeepLabCut projects.
| 31 | |
| 32 | |
| 33 | class DLCLoader(Loader): |
| 34 | """A Loader for DeepLabCut projects.""" |
| 35 | |
| 36 | def __init__( |
| 37 | self, |
| 38 | config: str | Path | dict, |
| 39 | trainset_index: int = 0, |
| 40 | shuffle: int = 0, |
| 41 | modelprefix: str = "", |
| 42 | ): |
| 43 | """ |
| 44 | Args: |
| 45 | config: Path to the DeepLabCut project config, or the project config itself |
| 46 | trainset_index: the index of the TrainingsetFraction for which to load data |
| 47 | shuffle: the index of the shuffle for which to load data |
| 48 | modelprefix: the modelprefix for the shuffle |
| 49 | """ |
| 50 | if isinstance(config, (str, Path)): |
| 51 | self._project_root = Path(config).parent |
| 52 | self._project_config = af.read_config(str(config)) |
| 53 | else: |
| 54 | self._project_root = Path(config["project_path"]) |
| 55 | self._project_config = config |
| 56 | |
| 57 | self._shuffle = shuffle |
| 58 | self._trainset_index = trainset_index |
| 59 | self._train_frac = self._project_config["TrainingFraction"][trainset_index] |
| 60 | self._model_folder = af.get_model_folder( |
| 61 | self._train_frac, |
| 62 | shuffle, |
| 63 | self._project_config, |
| 64 | engine=Engine.PYTORCH, |
| 65 | modelprefix=modelprefix, |
| 66 | ) |
| 67 | self._evaluation_folder = af.get_evaluation_folder( |
| 68 | trainFraction=self._train_frac, |
| 69 | shuffle=shuffle, |
| 70 | cfg=self._project_config, |
| 71 | engine=Engine.PYTORCH, |
| 72 | modelprefix=modelprefix, |
| 73 | ) |
| 74 | model_config_path = self._project_root / self._model_folder / "train" / Engine.PYTORCH.pose_cfg_name |
| 75 | super().__init__(self._project_root, self._project_root, model_config_path) |
| 76 | |
| 77 | # lazy-load split and DataFrames |
| 78 | self._split: dict[str, list[int]] | None = None |
| 79 | self._loaded_df: dict[str, pd.DataFrame] | None = None |
| 80 | self._resolutions = set() |
| 81 | |
| 82 | @property |
| 83 | def project_cfg(self) -> dict: |
| 84 | """Returns: the configuration for the DeepLabCut project""" |
| 85 | return self._project_config |
| 86 | |
| 87 | @property |
| 88 | def df(self) -> pd.DataFrame: |
| 89 | """Returns: The ground truth dataframe. Should not be modified.""" |
| 90 | return self._dfs["full"] |
no outgoing calls