deeplabcut.pose_estimation_pytorch.data.dlcloader

Class implementing the Loader for DeepLabCut projects.

Classes:

Name	Description
DLCLoader	A Loader for DeepLabCut projects.

Functions:

Name	Description
build_dlc_dataframe_columns	Builds the columns for a DeepLabCut DataFrame.

DLCLoader

Bases: Loader

A Loader for DeepLabCut projects.

Methods:

Name	Description
__init__	Args:
get_dataset_parameters	Retrieves dataset parameters based on the instance's configuration.
image_resolutions	Returns: The collection of image resolutions present in the dataset
load_data	Loads DeepLabCut data into COCO-style annotations.
load_ground_truth	Loads the ground truth dataset for a DeepLabCut project.
load_split	Loads the train/test split for a DeepLabCut shuffle.
scorer	Returns the scorer for this DLCLoader and the given snapshot.
split_data	Splits a DeepLabCut DataFrame into train/test dataframes.
to_coco	Formerly Shaokai's function.

Attributes:

Name	Type	Description
df	DataFrame	Returns: The ground truth dataframe. Should not be modified.
df_test	DataFrame	Returns: A copy of the DataFrame containing the test data.
df_train	DataFrame	Returns: A copy of the DataFrame containing the training data.
evaluation_folder	Path	Returns: The path to the evaluation folder
project_cfg	dict	Returns: the configuration for the DeepLabCut project
project_path	Path	Returns: The path to the DeepLabCut project
shuffle	int	Returns: the shuffle being loaded
train_fraction	float	Returns: the fraction of the dataset used for training

Source code in deeplabcut/pose_estimation_pytorch/data/dlcloader.py

class DLCLoader(Loader):
    """A Loader for DeepLabCut projects."""

    def __init__(
        self,
        config: str | Path | dict,
        trainset_index: int = 0,
        shuffle: int = 0,
        modelprefix: str = "",
    ):
        """
        Args:
            config: Path to the DeepLabCut project config, or the project config itself
            trainset_index: the index of the TrainingsetFraction for which to load data
            shuffle: the index of the shuffle for which to load data
            modelprefix: the modelprefix for the shuffle
        """
        if isinstance(config, (str, Path)):
            self._project_root = Path(config).parent
            self._project_config = af.read_config(str(config))
        else:
            self._project_root = Path(config["project_path"])
            self._project_config = config

        self._shuffle = shuffle
        self._trainset_index = trainset_index
        self._train_frac = self._project_config["TrainingFraction"][trainset_index]
        self._model_folder = af.get_model_folder(
            self._train_frac,
            shuffle,
            self._project_config,
            engine=Engine.PYTORCH,
            modelprefix=modelprefix,
        )
        self._evaluation_folder = af.get_evaluation_folder(
            trainFraction=self._train_frac,
            shuffle=shuffle,
            cfg=self._project_config,
            engine=Engine.PYTORCH,
            modelprefix=modelprefix,
        )
        model_config_path = self._project_root / self._model_folder / "train" / Engine.PYTORCH.pose_cfg_name
        super().__init__(self._project_root, self._project_root, model_config_path)

        # lazy-load split and DataFrames
        self._split: dict[str, list[int]] | None = None
        self._loaded_df: dict[str, pd.DataFrame] | None = None
        self._resolutions = set()

    @property
    def project_cfg(self) -> dict:
        """Returns: the configuration for the DeepLabCut project"""
        return self._project_config

    @property
    def df(self) -> pd.DataFrame:
        """Returns: The ground truth dataframe. Should not be modified."""
        return self._dfs["full"]

    @property
    def df_test(self) -> pd.DataFrame:
        """Returns: A copy of the DataFrame containing the test data."""
        return self._dfs["test"].copy()

    @property
    def df_train(self) -> pd.DataFrame:
        """Returns: A copy of the DataFrame containing the training data."""
        return self._dfs["train"].copy()

    def image_resolutions(self) -> set[tuple[int, int]]:
        """Returns: The collection of image resolutions present in the dataset"""
        return self._resolutions

    @property
    def evaluation_folder(self) -> Path:
        """Returns: The path to the evaluation folder"""
        return self._project_root / self._evaluation_folder

    @property
    def project_path(self) -> Path:
        """Returns: The path to the DeepLabCut project"""
        return self._project_root

    @property
    def shuffle(self) -> int:
        """Returns: the shuffle being loaded"""
        return self._shuffle

    @property
    def train_fraction(self) -> float:
        """Returns: the fraction of the dataset used for training"""
        return self._train_frac

    @property
    def split(self) -> dict[str, list[int]]:
        if self._split is None:
            self._split = self.load_split(self._project_config, self._trainset_index, self.shuffle)

        return self._split

    def scorer(
        self,
        snapshot: Snapshot | str | Path,
        detector_snapshot: Snapshot | str | Path | None = None,
    ) -> str:
        """Returns the scorer for this DLCLoader and the given snapshot."""
        task, date = self.project_cfg["Task"], self.project_cfg["date"]
        name = "".join([p.capitalize() for p in self.model_cfg["net_type"].split("_")])

        if not isinstance(snapshot, Snapshot):
            snapshot = Snapshot.from_path(Path(snapshot))

        snapshot_id = f"snapshot_{snapshot.uid()}"
        if detector_snapshot is not None:
            if not isinstance(detector_snapshot, Snapshot):
                detector_snapshot = Snapshot.from_path(Path(detector_snapshot))

            detect_id = detector_snapshot.uid()
            snapshot_id = f"detector_{detect_id}_{snapshot_id}"

        return f"DLC_{name}_{task}{date}shuffle{self.shuffle}_{snapshot_id}"

    def get_dataset_parameters(self) -> PoseDatasetParameters:
        """Retrieves dataset parameters based on the instance's configuration.

        Returns:
            An instance of the PoseDatasetParameters with the parameters set.
        """
        crop_cfg = self.model_cfg["data"]["train"].get("top_down_crop", {})
        crop_w, crop_h = crop_cfg.get("width", 256), crop_cfg.get("height", 256)
        crop_margin = crop_cfg.get("margin", 0)
        crop_with_context = crop_cfg.get("crop_with_context", True)

        return PoseDatasetParameters(
            bodyparts=self.model_cfg["metadata"]["bodyparts"],
            unique_bpts=self.model_cfg["metadata"]["unique_bodyparts"],
            individuals=self.model_cfg["metadata"]["individuals"],
            with_center_keypoints=self.model_cfg.get("with_center_keypoints", False),
            color_mode=self.model_cfg.get("color_mode", "RGB"),
            top_down_crop_size=(crop_w, crop_h),
            top_down_crop_margin=crop_margin,
            top_down_crop_with_context=crop_with_context,
        )

    def load_data(self, mode: str = "train") -> dict:
        """Loads DeepLabCut data into COCO-style annotations.

        This function reads data from h5 file, split the data and returns it in
        COCO-like format

        Args:
            mode: mode indicating whether to use 'train' or 'test' data.

        Raises:
            AttributeError: if the specified mode (train or test) does not exist.

        Returns:
            the coco-style annotations
        """
        if mode not in ["train", "test"]:
            raise AttributeError(f"Unknown mode: {mode}")
        if mode not in self._dfs:
            raise ValueError(f"No split for: {mode} (found {self._dfs.keys()})")
        if self._dfs[mode] is None:
            raise ValueError(f"No data in {mode} split for this shuffle!")

        params = self.get_dataset_parameters()
        data = self.to_coco(str(self._project_root), self._dfs[mode], params)
        with_bbox = self._compute_bboxes(
            data["images"],
            data["annotations"],
            method="keypoints",
            bbox_margin=self.model_cfg["data"].get("bbox_margin", 20),
        )
        data["annotations"] = with_bbox
        return data

    def load_ground_truth(
        self,
        config: dict,
        trainset_index: int,
        shuffle: int,
    ) -> tuple[dict[str, pd.DataFrame], set[tuple[int, int]]]:
        """Loads the ground truth dataset for a DeepLabCut project.

        Args:
            config: the DeepLabCut project configuration file
            trainset_index: the TrainingsetFraction for which to load data
            shuffle: the index of the shuffle for which to load data

        Returns: ground_truth_dataframes, image_resolutions
            ground_truth_dataframes: a dictionary containing the different DataFrames
                for the annotated DeepLabCut data for the current iteration
            image_resolutions: all possible image resolutions in the dataset

        Raises:
            ValueError: if the data contained in the ground truth HDF does not contain
                a dataframe.
        """
        trainset_dir = Path(config["project_path"]) / af.get_training_set_folder(config)
        dataset_path = f"CollectedData_{config['scorer']}.h5"
        train_frac = int(100 * config["TrainingFraction"][trainset_index])
        project_id = f"{config['Task']}_{config['scorer']}"
        dataset_file = trainset_dir / f"{project_id}{train_frac}shuffle{shuffle}"
        params = self.get_dataset_parameters()

        # as in TF DeepLabCut, load the training data from the .mat/.pickle file
        if config.get("multianimalproject", False):
            image_sizes, df_train = _load_pickle_dataset(
                dataset_file.with_suffix(".pickle"),
                config["scorer"],
                params=params,
            )
        else:
            image_sizes, df_train = _load_mat_dataset(
                dataset_file.with_suffix(".mat"),
                config["scorer"],
                params=params,
            )

        # load the full dataset file
        df = pd.read_hdf(trainset_dir / dataset_path)
        if not isinstance(df, pd.DataFrame):
            raise ValueError(f"The ground truth data in {trainset_dir} must contain a DataFrame! Found {df}")

        # load the data splits, check that there's nothing suspect
        dfs = self.split_data(df, self.split)
        dfs["full"] = df
        # let's not validate for now
        # dfs = _validate_dataframes(dfs, df_train)
        return dfs, image_sizes

    @staticmethod
    def load_split(
        config: dict,
        trainset_index: int = 0,
        shuffle: int = 0,
    ) -> dict[str, list[int]]:
        """Loads the train/test split for a DeepLabCut shuffle.

        Args:
            config: the DeepLabCut project config
            trainset_index: the TrainingsetFraction for which to load data
            shuffle: the index of the shuffle for which to load data

        Return:
            the {"train": [train_ids], "test": [test_ids]} data split
        """
        trainset_dir = Path(config["project_path"]) / af.get_training_set_folder(config)
        train_frac = int(100 * config["TrainingFraction"][trainset_index])
        shuffle_id = f"{config['Task']}_{train_frac}shuffle{shuffle}.pickle"
        doc_path = trainset_dir / f"Documentation_data-{shuffle_id}"

        with open(doc_path, "rb") as f:
            meta = pickle.load(f)

        train_ids = [int(i) for i in meta[1]]
        test_ids = [int(i) for i in meta[2]]
        return {"train": train_ids, "test": test_ids}

    @staticmethod
    def load_predictions(
        bu_snapshot: Path,
        bu_predictions: Path,
        parameters: PoseDatasetParameters,
    ) -> pd.DataFrame:
        if bu_predictions is None:
            pred_path = Path(str(bu_snapshot).replace("dlc-models", "evaluation-results")).parent.parent
            cfg = af.read_config(pred_path.parent.parent.parent / "config.yaml")
            scorer = af.get_scorer_name(
                cfg=cfg,
                shuffle=int(re.search(r"shuffle(\d+)", str(bu_snapshot)).group(1)),
                trainFraction=int(re.search(r"trainset(\d+)", str(bu_snapshot)).group(1)) / 100,
                engine=Engine.PYTORCH,
                trainingsiterations=re.search(r"snapshot-(.+)\.pth", str(bu_snapshot)).group(1),
                modelprefix="",
            )

            pred_file = pred_path / f"{scorer[0]}.h5"
            dlc_preds = pd.read_hdf(pred_file, key="df_with_missing")

            # FIXME: Implement the case where snapshot is loaded
            raise NotImplementedError("Need to implement the case with loaded snapshot")

        else:
            pred_path = bu_predictions.parent.parent
            dlc_preds = pd.read_hdf(bu_predictions, key="df_with_missing")

        predictions = {}
        for idx in dlc_preds.index.unique():
            if isinstance(idx, tuple):
                img_path = pred_path.parent.parent / Path(*idx)
            else:
                img_path = pred_path.parent.parent / Path(idx)

            keypoints = dlc_preds.loc[idx].values.reshape(-1, len(parameters.bodyparts), 3)[..., :2]
            keypoints = keypoints[~np.isnan(keypoints).all(axis=-1).all(axis=-1)]
            cond_keypoints = np.zeros((*keypoints.shape[:-1], 3))
            cond_keypoints[..., :2] = keypoints
            cond_keypoints[..., 2] = 2
            predictions[str(img_path)] = cond_keypoints

        return predictions

    @staticmethod
    def split_data(
        dlc_df: pd.DataFrame,
        split: dict[str, list[int]],
    ) -> dict[str, pd.DataFrame | None]:
        """Splits a DeepLabCut DataFrame into train/test dataframes.

        Args:
            dlc_df: the dataframe containing the labeled data
            split: the train/test indices

        Returns:
            a dictionary containing the same keys as the split dictionary, where the
            values are the rows of dlc_df with index in the split, or None if there are
            no indices in that split
        """
        split_dfs = {}
        for k, indices in split.items():
            if len(indices) == 0:
                split_dfs[k] = None
            else:
                split_dfs[k] = dlc_df.iloc[indices]
        return split_dfs

    @staticmethod
    def to_coco(
        project_root: str | Path,
        df: pd.DataFrame,
        parameters: PoseDatasetParameters,
    ) -> dict:
        """Formerly Shaokai's function.

        Args:
            project_root: the path to the project root
            df: the DLC-format annotation dataframe to convert to a COCO-format dict
            parameters: the parameters for pose estimation

        Returns:
            the coco format data
        """
        df = drop_likelihood_columns(df)

        with_individuals = "individuals" in df.columns.names
        if not with_individuals and (len(parameters.individuals) > 1 or len(parameters.unique_bpts) > 0):
            raise ValueError(
                "The DataFrame contains single-animal annotations (for a single, "
                "individual), but the parameters suggest this is a multi-animal project"
                f": {parameters} (with multiple individuals or unique bodyparts)"
            )

        categories = [
            {
                "id": 1,
                "name": "animals",
                "supercategory": "animal",
                "keypoints": parameters.bodyparts,
            },
        ]
        individuals = [idv for idv in parameters.individuals]
        if len(parameters.unique_bpts) > 0:
            individuals += ["single"]
            categories.append(
                {
                    "id": 2,
                    "name": "unique_bodypart",
                    "supercategory": "animal",
                    "keypoints": parameters.unique_bpts,
                }
            )

        anns, images = [], []
        base_path = Path(project_root)
        for idx, row in df.iterrows():
            image_id = len(images) + 1
            rel_path = Path(*idx) if isinstance(idx, tuple) else Path(str(idx))
            path = str(base_path / rel_path)
            _, height, width = read_image_shape_fast(path)
            images.append(
                {
                    "id": image_id,
                    "file_name": path,
                    "width": width,
                    "height": height,
                }
            )

            for idv_idx, idv in enumerate(individuals):
                category_id = 1
                individual_id = idv_idx
                if with_individuals:
                    if idv == "single":
                        category_id = 2
                        individual_id = -1
                    data = row.xs(idv, level="individuals")
                else:
                    data = row

                raw_keypoints = data.to_numpy().reshape((-1, 2))
                keypoints = np.zeros((len(raw_keypoints), 3))
                keypoints[:, :2] = raw_keypoints
                is_visible = np.logical_and(
                    ~pd.isnull(raw_keypoints).all(axis=1),
                    np.logical_and(
                        np.logical_and(
                            0 < keypoints[..., 0],
                            keypoints[..., 0] < width,
                        ),
                        np.logical_and(
                            0 < keypoints[..., 1],
                            keypoints[..., 1] < height,
                        ),
                    ),
                )
                keypoints[:, 2] = np.where(is_visible, 2, 0)
                num_keypoints = is_visible.sum()
                if num_keypoints > 0:
                    anns.append(
                        {
                            "id": len(anns) + 1,
                            "image_id": image_id,
                            "category_id": category_id,
                            "individual": idv,
                            "individual_id": individual_id,
                            "num_keypoints": num_keypoints,
                            "keypoints": keypoints,
                            "iscrowd": 0,
                        }
                    )

        coco_dict = {"annotations": anns, "categories": categories, "images": images}
        coco_dict = DLCLoader._add_bbox_annotations(coco_dict)
        coco_dict = DLCLoader._remove_nans(coco_dict)
        return coco_dict

    @staticmethod
    def _add_bbox_annotations(coco_dict: dict) -> dict:
        for annotation in coco_dict.get("annotations", []):
            if "bbox" not in annotation:
                image = [img for img in coco_dict.get("images") if img.get("id") == annotation.get("image_id")][0]
                bbox = bbox_from_keypoints(
                    keypoints=np.array(annotation["keypoints"]),  # (..., num_keypoints, xy)
                    image_h=image.get("height"),
                    image_w=image.get("width"),
                    margin=20,
                )
                annotation["bbox"] = list(bbox)
        return coco_dict

    @staticmethod
    def _remove_nans(coco_dict: dict) -> dict:
        # Iterate through annotations and fix keypoints
        for annotation in coco_dict.get("annotations", []):
            if "keypoints" in annotation:
                for keypoint in annotation["keypoints"]:
                    if any(isinstance(v, float) and np.isnan(v) for v in keypoint[:2]):
                        keypoint[0] = 0.0  # Replace x with 0
                        keypoint[1] = 0.0  # Replace y with 0
                        keypoint[2] = 0.0  # Ensure visibility is also 0
        return coco_dict

    @property
    def _dfs(self) -> dict[str, pd.DataFrame]:
        """Lazy-loading of the training dataset dataframes."""
        if self._loaded_df is None:
            self._loaded_df, image_sizes = self.load_ground_truth(
                self._project_config,
                trainset_index=self._trainset_index,
                shuffle=self.shuffle,
            )
            self._resolutions = self._resolutions.union(image_sizes)

        return self._loaded_df

df property

df: DataFrame

Returns: The ground truth dataframe. Should not be modified.

df_test property

df_test: DataFrame

Returns: A copy of the DataFrame containing the test data.

df_train property

df_train: DataFrame

Returns: A copy of the DataFrame containing the training data.

evaluation_folder property

evaluation_folder: Path

Returns: The path to the evaluation folder

project_cfg property

project_cfg: dict

Returns: the configuration for the DeepLabCut project

project_path property

project_path: Path

Returns: The path to the DeepLabCut project

shuffle property

shuffle: int

Returns: the shuffle being loaded

train_fraction property

train_fraction: float

Returns: the fraction of the dataset used for training

__init__

__init__(config: str | Path | dict, trainset_index: int = 0, shuffle: int = 0, modelprefix: str = '')

Parameters:

Name	Type	Description	Default
config	str | Path | dict	Path to the DeepLabCut project config, or the project config itself	required
trainset_index	int	the index of the TrainingsetFraction for which to load data	0
shuffle	int	the index of the shuffle for which to load data	0
modelprefix	str	the modelprefix for the shuffle	''

Source code in deeplabcut/pose_estimation_pytorch/data/dlcloader.py

def __init__(
    self,
    config: str | Path | dict,
    trainset_index: int = 0,
    shuffle: int = 0,
    modelprefix: str = "",
):
    """
    Args:
        config: Path to the DeepLabCut project config, or the project config itself
        trainset_index: the index of the TrainingsetFraction for which to load data
        shuffle: the index of the shuffle for which to load data
        modelprefix: the modelprefix for the shuffle
    """
    if isinstance(config, (str, Path)):
        self._project_root = Path(config).parent
        self._project_config = af.read_config(str(config))
    else:
        self._project_root = Path(config["project_path"])
        self._project_config = config

    self._shuffle = shuffle
    self._trainset_index = trainset_index
    self._train_frac = self._project_config["TrainingFraction"][trainset_index]
    self._model_folder = af.get_model_folder(
        self._train_frac,
        shuffle,
        self._project_config,
        engine=Engine.PYTORCH,
        modelprefix=modelprefix,
    )
    self._evaluation_folder = af.get_evaluation_folder(
        trainFraction=self._train_frac,
        shuffle=shuffle,
        cfg=self._project_config,
        engine=Engine.PYTORCH,
        modelprefix=modelprefix,
    )
    model_config_path = self._project_root / self._model_folder / "train" / Engine.PYTORCH.pose_cfg_name
    super().__init__(self._project_root, self._project_root, model_config_path)

    # lazy-load split and DataFrames
    self._split: dict[str, list[int]] | None = None
    self._loaded_df: dict[str, pd.DataFrame] | None = None
    self._resolutions = set()

get_dataset_parameters

get_dataset_parameters() -> PoseDatasetParameters

Retrieves dataset parameters based on the instance's configuration.

Returns:

Type	Description
PoseDatasetParameters	An instance of the PoseDatasetParameters with the parameters set.

Source code in deeplabcut/pose_estimation_pytorch/data/dlcloader.py

def get_dataset_parameters(self) -> PoseDatasetParameters:
    """Retrieves dataset parameters based on the instance's configuration.

    Returns:
        An instance of the PoseDatasetParameters with the parameters set.
    """
    crop_cfg = self.model_cfg["data"]["train"].get("top_down_crop", {})
    crop_w, crop_h = crop_cfg.get("width", 256), crop_cfg.get("height", 256)
    crop_margin = crop_cfg.get("margin", 0)
    crop_with_context = crop_cfg.get("crop_with_context", True)

    return PoseDatasetParameters(
        bodyparts=self.model_cfg["metadata"]["bodyparts"],
        unique_bpts=self.model_cfg["metadata"]["unique_bodyparts"],
        individuals=self.model_cfg["metadata"]["individuals"],
        with_center_keypoints=self.model_cfg.get("with_center_keypoints", False),
        color_mode=self.model_cfg.get("color_mode", "RGB"),
        top_down_crop_size=(crop_w, crop_h),
        top_down_crop_margin=crop_margin,
        top_down_crop_with_context=crop_with_context,
    )

image_resolutions

image_resolutions() -> set[tuple[int, int]]

Returns: The collection of image resolutions present in the dataset

Source code in deeplabcut/pose_estimation_pytorch/data/dlcloader.py

def image_resolutions(self) -> set[tuple[int, int]]:
    """Returns: The collection of image resolutions present in the dataset"""
    return self._resolutions

load_data

load_data(mode: str = 'train') -> dict

Loads DeepLabCut data into COCO-style annotations.

This function reads data from h5 file, split the data and returns it in COCO-like format

Parameters:

Name	Type	Description	Default
mode	str	mode indicating whether to use 'train' or 'test' data.	'train'

Raises:

Type	Description
AttributeError	if the specified mode (train or test) does not exist.

Returns:

Type	Description
dict	the coco-style annotations

Source code in deeplabcut/pose_estimation_pytorch/data/dlcloader.py

def load_data(self, mode: str = "train") -> dict:
    """Loads DeepLabCut data into COCO-style annotations.

    This function reads data from h5 file, split the data and returns it in
    COCO-like format

    Args:
        mode: mode indicating whether to use 'train' or 'test' data.

    Raises:
        AttributeError: if the specified mode (train or test) does not exist.

    Returns:
        the coco-style annotations
    """
    if mode not in ["train", "test"]:
        raise AttributeError(f"Unknown mode: {mode}")
    if mode not in self._dfs:
        raise ValueError(f"No split for: {mode} (found {self._dfs.keys()})")
    if self._dfs[mode] is None:
        raise ValueError(f"No data in {mode} split for this shuffle!")

    params = self.get_dataset_parameters()
    data = self.to_coco(str(self._project_root), self._dfs[mode], params)
    with_bbox = self._compute_bboxes(
        data["images"],
        data["annotations"],
        method="keypoints",
        bbox_margin=self.model_cfg["data"].get("bbox_margin", 20),
    )
    data["annotations"] = with_bbox
    return data

load_ground_truth

load_ground_truth(
    config: dict, trainset_index: int, shuffle: int
) -> tuple[dict[str, pd.DataFrame], set[tuple[int, int]]]

Loads the ground truth dataset for a DeepLabCut project.

Parameters:

Name	Type	Description	Default
config	dict	the DeepLabCut project configuration file	required
trainset_index	int	the TrainingsetFraction for which to load data	required
shuffle	int	the index of the shuffle for which to load data	required

ground_truth_dataframes, image_resolutions

Name	Type	Description
ground_truth_dataframes	tuple[dict[str, DataFrame], set[tuple[int, int]]]	a dictionary containing the different DataFrames for the annotated DeepLabCut data for the current iteration image_resolutions: all possible image resolutions in the dataset

Raises:

Type	Description
ValueError	if the data contained in the ground truth HDF does not contain a dataframe.

Source code in deeplabcut/pose_estimation_pytorch/data/dlcloader.py

def load_ground_truth(
    self,
    config: dict,
    trainset_index: int,
    shuffle: int,
) -> tuple[dict[str, pd.DataFrame], set[tuple[int, int]]]:
    """Loads the ground truth dataset for a DeepLabCut project.

    Args:
        config: the DeepLabCut project configuration file
        trainset_index: the TrainingsetFraction for which to load data
        shuffle: the index of the shuffle for which to load data

    Returns: ground_truth_dataframes, image_resolutions
        ground_truth_dataframes: a dictionary containing the different DataFrames
            for the annotated DeepLabCut data for the current iteration
        image_resolutions: all possible image resolutions in the dataset

    Raises:
        ValueError: if the data contained in the ground truth HDF does not contain
            a dataframe.
    """
    trainset_dir = Path(config["project_path"]) / af.get_training_set_folder(config)
    dataset_path = f"CollectedData_{config['scorer']}.h5"
    train_frac = int(100 * config["TrainingFraction"][trainset_index])
    project_id = f"{config['Task']}_{config['scorer']}"
    dataset_file = trainset_dir / f"{project_id}{train_frac}shuffle{shuffle}"
    params = self.get_dataset_parameters()

    # as in TF DeepLabCut, load the training data from the .mat/.pickle file
    if config.get("multianimalproject", False):
        image_sizes, df_train = _load_pickle_dataset(
            dataset_file.with_suffix(".pickle"),
            config["scorer"],
            params=params,
        )
    else:
        image_sizes, df_train = _load_mat_dataset(
            dataset_file.with_suffix(".mat"),
            config["scorer"],
            params=params,
        )

    # load the full dataset file
    df = pd.read_hdf(trainset_dir / dataset_path)
    if not isinstance(df, pd.DataFrame):
        raise ValueError(f"The ground truth data in {trainset_dir} must contain a DataFrame! Found {df}")

    # load the data splits, check that there's nothing suspect
    dfs = self.split_data(df, self.split)
    dfs["full"] = df
    # let's not validate for now
    # dfs = _validate_dataframes(dfs, df_train)
    return dfs, image_sizes

load_split staticmethod

load_split(config: dict, trainset_index: int = 0, shuffle: int = 0) -> dict[str, list[int]]

Loads the train/test split for a DeepLabCut shuffle.

Parameters:

Name	Type	Description	Default
config	dict	the DeepLabCut project config	required
trainset_index	int	the TrainingsetFraction for which to load data	0
shuffle	int	the index of the shuffle for which to load data	0

Return

the {"train": [train_ids], "test": [test_ids]} data split

Source code in deeplabcut/pose_estimation_pytorch/data/dlcloader.py

@staticmethod
def load_split(
    config: dict,
    trainset_index: int = 0,
    shuffle: int = 0,
) -> dict[str, list[int]]:
    """Loads the train/test split for a DeepLabCut shuffle.

    Args:
        config: the DeepLabCut project config
        trainset_index: the TrainingsetFraction for which to load data
        shuffle: the index of the shuffle for which to load data

    Return:
        the {"train": [train_ids], "test": [test_ids]} data split
    """
    trainset_dir = Path(config["project_path"]) / af.get_training_set_folder(config)
    train_frac = int(100 * config["TrainingFraction"][trainset_index])
    shuffle_id = f"{config['Task']}_{train_frac}shuffle{shuffle}.pickle"
    doc_path = trainset_dir / f"Documentation_data-{shuffle_id}"

    with open(doc_path, "rb") as f:
        meta = pickle.load(f)

    train_ids = [int(i) for i in meta[1]]
    test_ids = [int(i) for i in meta[2]]
    return {"train": train_ids, "test": test_ids}

scorer

scorer(snapshot: Snapshot | str | Path, detector_snapshot: Snapshot | str | Path | None = None) -> str

Returns the scorer for this DLCLoader and the given snapshot.

Source code in deeplabcut/pose_estimation_pytorch/data/dlcloader.py

def scorer(
    self,
    snapshot: Snapshot | str | Path,
    detector_snapshot: Snapshot | str | Path | None = None,
) -> str:
    """Returns the scorer for this DLCLoader and the given snapshot."""
    task, date = self.project_cfg["Task"], self.project_cfg["date"]
    name = "".join([p.capitalize() for p in self.model_cfg["net_type"].split("_")])

    if not isinstance(snapshot, Snapshot):
        snapshot = Snapshot.from_path(Path(snapshot))

    snapshot_id = f"snapshot_{snapshot.uid()}"
    if detector_snapshot is not None:
        if not isinstance(detector_snapshot, Snapshot):
            detector_snapshot = Snapshot.from_path(Path(detector_snapshot))

        detect_id = detector_snapshot.uid()
        snapshot_id = f"detector_{detect_id}_{snapshot_id}"

    return f"DLC_{name}_{task}{date}shuffle{self.shuffle}_{snapshot_id}"

split_data staticmethod

split_data(dlc_df: DataFrame, split: dict[str, list[int]]) -> dict[str, pd.DataFrame | None]

Splits a DeepLabCut DataFrame into train/test dataframes.

Parameters:

Name	Type	Description	Default
dlc_df	DataFrame	the dataframe containing the labeled data	required
split	dict[str, list[int]]	the train/test indices	required

Returns:

Type	Description
dict[str, DataFrame | None]	a dictionary containing the same keys as the split dictionary, where the values are the rows of dlc_df with index in the split, or None if there are no indices in that split

Source code in deeplabcut/pose_estimation_pytorch/data/dlcloader.py

@staticmethod
def split_data(
    dlc_df: pd.DataFrame,
    split: dict[str, list[int]],
) -> dict[str, pd.DataFrame | None]:
    """Splits a DeepLabCut DataFrame into train/test dataframes.

    Args:
        dlc_df: the dataframe containing the labeled data
        split: the train/test indices

    Returns:
        a dictionary containing the same keys as the split dictionary, where the
        values are the rows of dlc_df with index in the split, or None if there are
        no indices in that split
    """
    split_dfs = {}
    for k, indices in split.items():
        if len(indices) == 0:
            split_dfs[k] = None
        else:
            split_dfs[k] = dlc_df.iloc[indices]
    return split_dfs

to_coco staticmethod

to_coco(project_root: str | Path, df: DataFrame, parameters: PoseDatasetParameters) -> dict

Formerly Shaokai's function.

Parameters:

Name	Type	Description	Default
project_root	str | Path	the path to the project root	required
df	DataFrame	the DLC-format annotation dataframe to convert to a COCO-format dict	required
parameters	PoseDatasetParameters	the parameters for pose estimation	required

Returns:

Type	Description
dict	the coco format data

Source code in deeplabcut/pose_estimation_pytorch/data/dlcloader.py

@staticmethod
def to_coco(
    project_root: str | Path,
    df: pd.DataFrame,
    parameters: PoseDatasetParameters,
) -> dict:
    """Formerly Shaokai's function.

    Args:
        project_root: the path to the project root
        df: the DLC-format annotation dataframe to convert to a COCO-format dict
        parameters: the parameters for pose estimation

    Returns:
        the coco format data
    """
    df = drop_likelihood_columns(df)

    with_individuals = "individuals" in df.columns.names
    if not with_individuals and (len(parameters.individuals) > 1 or len(parameters.unique_bpts) > 0):
        raise ValueError(
            "The DataFrame contains single-animal annotations (for a single, "
            "individual), but the parameters suggest this is a multi-animal project"
            f": {parameters} (with multiple individuals or unique bodyparts)"
        )

    categories = [
        {
            "id": 1,
            "name": "animals",
            "supercategory": "animal",
            "keypoints": parameters.bodyparts,
        },
    ]
    individuals = [idv for idv in parameters.individuals]
    if len(parameters.unique_bpts) > 0:
        individuals += ["single"]
        categories.append(
            {
                "id": 2,
                "name": "unique_bodypart",
                "supercategory": "animal",
                "keypoints": parameters.unique_bpts,
            }
        )

    anns, images = [], []
    base_path = Path(project_root)
    for idx, row in df.iterrows():
        image_id = len(images) + 1
        rel_path = Path(*idx) if isinstance(idx, tuple) else Path(str(idx))
        path = str(base_path / rel_path)
        _, height, width = read_image_shape_fast(path)
        images.append(
            {
                "id": image_id,
                "file_name": path,
                "width": width,
                "height": height,
            }
        )

        for idv_idx, idv in enumerate(individuals):
            category_id = 1
            individual_id = idv_idx
            if with_individuals:
                if idv == "single":
                    category_id = 2
                    individual_id = -1
                data = row.xs(idv, level="individuals")
            else:
                data = row

            raw_keypoints = data.to_numpy().reshape((-1, 2))
            keypoints = np.zeros((len(raw_keypoints), 3))
            keypoints[:, :2] = raw_keypoints
            is_visible = np.logical_and(
                ~pd.isnull(raw_keypoints).all(axis=1),
                np.logical_and(
                    np.logical_and(
                        0 < keypoints[..., 0],
                        keypoints[..., 0] < width,
                    ),
                    np.logical_and(
                        0 < keypoints[..., 1],
                        keypoints[..., 1] < height,
                    ),
                ),
            )
            keypoints[:, 2] = np.where(is_visible, 2, 0)
            num_keypoints = is_visible.sum()
            if num_keypoints > 0:
                anns.append(
                    {
                        "id": len(anns) + 1,
                        "image_id": image_id,
                        "category_id": category_id,
                        "individual": idv,
                        "individual_id": individual_id,
                        "num_keypoints": num_keypoints,
                        "keypoints": keypoints,
                        "iscrowd": 0,
                    }
                )

    coco_dict = {"annotations": anns, "categories": categories, "images": images}
    coco_dict = DLCLoader._add_bbox_annotations(coco_dict)
    coco_dict = DLCLoader._remove_nans(coco_dict)
    return coco_dict

build_dlc_dataframe_columns

build_dlc_dataframe_columns(scorer: str, parameters: PoseDatasetParameters, with_likelihood: bool) -> pd.MultiIndex

Builds the columns for a DeepLabCut DataFrame.

Parameters:

Name	Type	Description	Default
scorer	str	the scorer name	required
parameters	PoseDatasetParameters	the parameters for the project	required
with_likelihood	bool	whether the DataFrame contains pose likelihood	required

Returns:

Type	Description
MultiIndex	the multi-index columns for the DataFrame

Source code in deeplabcut/pose_estimation_pytorch/data/dlcloader.py

def build_dlc_dataframe_columns(
    scorer: str,
    parameters: PoseDatasetParameters,
    with_likelihood: bool,
) -> pd.MultiIndex:
    """Builds the columns for a DeepLabCut DataFrame.

    Args:
        scorer: the scorer name
        parameters: the parameters for the project
        with_likelihood: whether the DataFrame contains pose likelihood

    Returns:
        the multi-index columns for the DataFrame
    """
    levels = ["scorer", "individuals", "bodyparts", "coords"]
    kpt_entries = ["x", "y"]
    if with_likelihood:
        kpt_entries.append("likelihood")

    columns = []
    for i in parameters.individuals:
        for b in parameters.bodyparts:
            columns += [(scorer, i, b, entry) for entry in kpt_entries]

    for unique_bpt in parameters.unique_bpts:
        columns += [(scorer, "single", unique_bpt, entry) for entry in kpt_entries]

    return pd.MultiIndex.from_tuples(columns, names=levels)