dlc_reader.py

PoseEstimation

Source code in src/spyglass/position/v1/dlc_reader.py

class PoseEstimation:
    def __init__(
        self,
        dlc_dir=None,
        pkl_path=None,
        h5_path=None,
        yml_path=None,
        filename_prefix="",
    ):
        ActivityLog.deprecate_log("dlc_reader: PoseEstimation")
        if dlc_dir is None:
            assert pkl_path and h5_path and yml_path, (
                'If "dlc_dir" is not provided, then pkl_path, h5_path, and yml_path '
                + "must be provided"
            )
        else:
            self.dlc_dir = Path(dlc_dir)
            assert self.dlc_dir.exists(), f"Unable to find {dlc_dir}"

        # meta file: pkl - info about this  DLC run (input video, configuration, etc.)
        if pkl_path is None:
            pkl_paths = list(
                self.dlc_dir.rglob(f"{filename_prefix}*meta.pickle")
            )
            if not test_mode:
                assert len(pkl_paths) == 1, (
                    "Unable to find one unique .pickle file in: "
                    + f"{dlc_dir} - Found: {len(pkl_paths)}"
                )
            self.pkl_path = pkl_paths[0]
        else:
            self.pkl_path = Path(pkl_path)
            assert self.pkl_path.exists()

        # data file: h5 - body part outputs from the DLC post estimation step
        if h5_path is None:
            h5_paths = list(self.dlc_dir.rglob(f"{filename_prefix}*.h5"))
            if not test_mode:
                assert len(h5_paths) == 1, (
                    "Unable to find one unique .h5 file in: "
                    + f"{dlc_dir} - Found: {len(h5_paths)}"
                )
            self.h5_path = h5_paths[0]
        else:
            self.h5_path = Path(h5_path)
            assert self.h5_path.exists()

        if not test_mode:
            assert (
                self.pkl_path.stem == self.h5_path.stem + "_meta"
            ), f"Mismatching h5 ({self.h5_path.stem}) and pickle {self.pkl_path.stem}"

        # config file: yaml - configuration for invoking the DLC post estimation step
        if yml_path is None:
            yml_paths = list(self.dlc_dir.glob(f"{filename_prefix}*.y*ml"))
            # If multiple, defer to the one we save.
            if len(yml_paths) > 1:
                yml_paths = [
                    val for val in yml_paths if val.stem == "dj_dlc_config"
                ]
            if not test_mode:
                assert len(yml_paths) == 1, (
                    "Unable to find one unique .yaml file in: "
                    + f"{dlc_dir} - Found: {len(yml_paths)}"
                )
            self.yml_path = yml_paths[0]
        else:
            self.yml_path = Path(yml_path)
            assert self.yml_path.exists()

        self._pkl = None
        self._rawdata = None
        self._yml = None
        self._data = None

        train_idx = np.where(
            (np.array(self.yml["TrainingFraction"]) * 100).astype(int)
            == int(self.pkl["training set fraction"] * 100)
        )[0][0]
        train_iter = int(self.pkl["Scorer"].split("_")[-1])

        self.model = {
            "Scorer": self.pkl["Scorer"],
            "Task": self.yml["Task"],
            "date": self.yml["date"],
            "iteration": self.pkl["iteration (active-learning)"],
            "shuffle": int(
                re.search(r"shuffle(\d+)", self.pkl["Scorer"]).groups()[0]
            ),
            "snapshotindex": self.yml["snapshotindex"],
            "trainingsetindex": train_idx,
            "training_iteration": train_iter,
        }

        self.fps = self.pkl["fps"]
        self.nframes = self.pkl["nframes"]

        self.creation_time = self.h5_path.stat().st_mtime

    @property
    def pkl(self):
        """Pickle object with metadata about the DLC run."""
        if self._pkl is None:
            with open(self.pkl_path, "rb") as f:
                self._pkl = pickle.load(f)
        return self._pkl["data"]

    @property  # DLC aux_func has a read_config option, but it rewrites the proj path
    def yml(self) -> dict:
        """Dictionary of the yaml file DLC metadata."""
        if self._yml is None:
            with open(self.yml_path, "rb") as f:
                safe_yaml = yaml.YAML(typ="safe", pure=True)
                self._yml = safe_yaml.load(f)
        return self._yml

    @property
    def rawdata(self):
        """Pandas dataframe of the DLC output from the h5 file."""
        if self._rawdata is None:
            self._rawdata = pd.read_hdf(self.h5_path)
        return self._rawdata

    @property
    def data(self) -> dict:
        """Dictionary of the bodyparts and corresponding dataframe data."""
        if self._data is None:
            self._data = self.reformat_rawdata()
        return self._data

    @property
    def df(self) -> pd.DataFrame:
        """Pandas dataframe of the DLC output from the h5 file."""
        top_level = self.rawdata.columns.levels[0][0]
        return self.rawdata.get(top_level)

    @property
    def body_parts(self) -> list[str]:
        """List of body parts in the DLC output."""
        return self.df.columns.levels[0]

    def reformat_rawdata(self) -> dict:
        """Reformat the rawdata from the h5 file to a more useful dictionary."""
        error_message = (
            f"Total frames from .h5 file ({len(self.rawdata)}) differs "
            + f'from .pickle ({self.pkl["nframes"]})'
        )
        assert len(self.rawdata) == self.pkl["nframes"], error_message

        body_parts_position = {}
        for body_part in self.body_parts:
            body_parts_position[body_part] = {
                c: self.df.get(body_part).get(c).values
                for c in self.df.get(body_part).columns
            }

        return body_parts_position

pkl property

Pickle object with metadata about the DLC run.

yml: dict property

Dictionary of the yaml file DLC metadata.

rawdata property

Pandas dataframe of the DLC output from the h5 file.

data: dict property

Dictionary of the bodyparts and corresponding dataframe data.

df: pd.DataFrame property

Pandas dataframe of the DLC output from the h5 file.

body_parts: list[str] property

List of body parts in the DLC output.

reformat_rawdata()

Reformat the rawdata from the h5 file to a more useful dictionary.

Source code in src/spyglass/position/v1/dlc_reader.py

def reformat_rawdata(self) -> dict:
    """Reformat the rawdata from the h5 file to a more useful dictionary."""
    error_message = (
        f"Total frames from .h5 file ({len(self.rawdata)}) differs "
        + f'from .pickle ({self.pkl["nframes"]})'
    )
    assert len(self.rawdata) == self.pkl["nframes"], error_message

    body_parts_position = {}
    for body_part in self.body_parts:
        body_parts_position[body_part] = {
            c: self.df.get(body_part).get(c).values
            for c in self.df.get(body_part).columns
        }

    return body_parts_position

read_yaml(fullpath, filename='*')

Return contents of yml in fullpath. If available, defer to DJ-saved version

Parameters:

Name	Type	Description	Default
fullpath		Directory with yaml files	required
filename		Filename, no extension. Permits wildcards.	'*'

Returns:

Type	Description
tuple	filepath and contents as dict

Source code in src/spyglass/position/v1/dlc_reader.py

def read_yaml(fullpath, filename="*"):
    """Return contents of yml in fullpath. If available, defer to DJ-saved version

    Parameters
    ----------
    fullpath: Union[str, pathlib.Path]
        Directory with yaml files
    filename: str
        Filename, no extension. Permits wildcards.

    Returns
    -------
    tuple
        filepath and contents as dict
    """
    from deeplabcut.utils.auxiliaryfunctions import read_config

    # Take the DJ-saved if there. If not, return list of available
    yml_paths = list(Path(fullpath).glob("dj_dlc_config.yaml")) or sorted(
        list(Path(fullpath).glob(f"{filename}.y*ml"))
    )

    assert (  # If more than 1 and not DJ-saved,
        len(yml_paths) == 1
    ), f"Found more yaml files than expected: {len(yml_paths)}\n{fullpath}"

    return yml_paths[0], read_config(yml_paths[0])

save_yaml(output_dir, config_dict, filename='dj_dlc_config', mkdir=True)

Save config_dict to output_path as filename.yaml. By default, preserves original.

Parameters:

Name	Type	Description	Default
output_dir			required
config_dict			required
filename		Set to 'config' to overwrite original file. If extension is included, removed and replaced with "yaml".	'dj_dlc_config'
mkdir			True

Returns:

Type	Description
str	path of saved file as string - due to DLC func preference for strings

Source code in src/spyglass/position/v1/dlc_reader.py

def save_yaml(output_dir, config_dict, filename="dj_dlc_config", mkdir=True):
    """Save config_dict to output_path as filename.yaml. By default, preserves original.

    Parameters
    ----------
    output_dir: where to save yaml file
    config_dict: dict of config params or element-deeplabcut model.Model dict
    filename: Optional, default 'dj_dlc_config' or preserve original 'config'
              Set to 'config' to overwrite original file.
              If extension is included, removed and replaced with "yaml".
    mkdir (bool): Optional, True. Make new directory if output_dir not exist

    Returns
    -------
    str
        path of saved file as string - due to DLC func preference for strings
    """
    from deeplabcut.utils.auxiliaryfunctions import write_config

    if "config_template" in config_dict:  # if passed full model.Model dict
        config_dict = config_dict["config_template"]
    if mkdir:
        Path(output_dir).mkdir(exist_ok=True)
    if "." in filename:  # if user provided extension, remove
        filename = filename.split(".")[0]

    output_filepath = Path(output_dir) / f"{filename}.yaml"
    write_config(output_filepath, config_dict)
    return str(output_filepath)

do_pose_estimation(video_filepaths, dlc_model, project_path, output_dir, videotype='', gputouse=None, save_as_csv=False, batchsize=None, cropping=None, TFGPUinference=True, dynamic=(False, 0.5, 10), robust_nframes=False, allow_growth=False, use_shelve=False)

Launch DLC's analyze_videos within element-deeplabcut

Other optional parameters may be set other than those described below. See deeplabcut.analyze_videos parameters for descriptions/defaults.

Parameters:

Name	Type	Description	Default
video_filepaths			required
dlc_model			required
project_path			required
output_dir			required

Source code in src/spyglass/position/v1/dlc_reader.py

def do_pose_estimation(
    video_filepaths,
    dlc_model,
    project_path,
    output_dir,
    videotype="",
    gputouse=None,
    save_as_csv=False,
    batchsize=None,
    cropping=None,
    TFGPUinference=True,
    dynamic=(False, 0.5, 10),
    robust_nframes=False,
    allow_growth=False,
    use_shelve=False,
):
    """Launch DLC's analyze_videos within element-deeplabcut

    Other optional parameters may be set other than those described below. See
    deeplabcut.analyze_videos parameters for descriptions/defaults.

    Parameters
    ----------
    video_filepaths: list of videos to analyze
    dlc_model: element-deeplabcut dlc.Model dict
    project_path: path to project config.yml
    output_dir: where to save output
    """
    from deeplabcut.pose_estimation_tensorflow import analyze_videos

    # ---- Build and save DLC configuration (yaml) file ----
    dlc_config = dlc_model["config_template"]
    dlc_project_path = Path(project_path)
    dlc_config["project_path"] = dlc_project_path.as_posix()

    # ---- Write config files ----
    # To output dir: Important for loading/parsing output in datajoint
    _ = save_yaml(output_dir, dlc_config)
    # To project dir: Required by DLC to run the analyze_videos
    if dlc_project_path != output_dir:
        config_filepath = save_yaml(dlc_project_path, dlc_config)
    # ---- Trigger DLC prediction job ----
    analyze_videos(
        config=config_filepath,
        videos=video_filepaths,
        shuffle=dlc_model["shuffle"],
        trainingsetindex=dlc_model["trainingsetindex"],
        destfolder=output_dir,
        modelprefix=dlc_model["model_prefix"],
        videotype=videotype,
        gputouse=gputouse,
        save_as_csv=save_as_csv,
        batchsize=batchsize,
        cropping=cropping,
        TFGPUinference=TFGPUinference,
        dynamic=dynamic,
        robust_nframes=robust_nframes,
        allow_growth=allow_growth,
        use_shelve=use_shelve,
    )