Visitar URL original
Fix shelf loading error by Ayskin · Pull Request #3558 · DeepLabCut/DeepLabCut · GitHub
Skip to content

Fix shelf loading error - #3558

Open
Ayskin wants to merge 9 commits into
DeepLabCut:devfrom
Ayskin:yann/fix-shelving-load
Open

Ayskin wants to merge 9 commits into
DeepLabCut:devfrom
Ayskin:yann/fix-shelving-load

Conversation

@Ayskin

@Ayskin Ayskin commented Oct 19, 2026 •

Copy link
Copy Markdown
Collaborator

Bug description

When running analyze_videos inference with use_shelve=True, the predictions holding file _full.pickle is written as a dbm database. Depending on the Python version and the available dbm backend (dbm.gnu, dbm.sqlite3, dbm.ndbm, etc), it can end up as different file names "_full.pickle", "_full.pickle.db", "_full.pickle.dir", etc.

Currently, the code only checks if the filename "_full.pickle" exists before calling pickle.load. This fails if "_full.pickle" (with no additional extension) is a shelf.

Fix

Centralized the predictions loading in deeplabcut/utils/auxiliaryfunctions.py:

  • open_full_data(path): context manager yielding a dict for a plain pickle, or a read-only shelf otherwise.
  • load_full_data(path): same, but returns an in-memory dict copy.
  • full_data_format(path): reads the type of the predictions file (None = no file, "" = plain pickle, else the dbm module name).
  • remove_full_data(path): removes existing files in case of overwrite.

These helpers fix the loading error that may happen at several places in the code : _generate_assemblies_file, CTD tracking, tracklets.py, outlier_frames.py, scoring.py, LoadFullMultiAnimalData.

Additional fix : overwrite=True is now safe to be used in analyze_videos as the helpers look for previously existing files and remove/overwrite them accordingly depending on their format.
This handles the issue where switching writing type (shelving or not) could leave coexisting files and mess up reading freshly created files.

Tests

  • Standard i/o tests for the helpers in tests/test_auxiliaryfunctions.py.

  • New tests to lock the fixes of analyze_videos in tests/pose_estimation_pytorch/apis/test_videos.py.

Tests parametrize on available dbm backends with unavailable ones being skipped.
I am not sure about best practice here so any comment on this is much appreciated.

@deruyter92 deruyter92 left a comment •

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@Ayskin great effort in this PR, good that you are addressing this!

A suggestion: what do you think about centralizing all of this logic into a single place, such that it can be reused at different call sites? What I currently see repeated is: what format are we using? Does the file exist on disk? Should we overwrite existing? This also matters for correctness. With the implemented backends, a shelf is stored as side files (e.g. .db) next to _full.pickle. If a user ran once with use_shelve=True and later without, both a fresh pickle and an old shelf exist, and dbm.whichdb reports the shelf.

Ideas for the centralized code:

  • the storage format could be persisted in _meta.pickle (e.g. "storage": "pickle" | "shelve").
  • a centralized context manager could return the same mapping regardless whether the file is a pickle or a shelf (or maybe a zarr in the future). The if/else branches can be dropped in all call sites (i.e. assemblies, CTD, LoadFullMultiAnimalData, outlier frames).
  • the writer owns the file: when an analysis actually (re)writes the output, it can remove leftovers of the other format (or refuse to open a shelf when a pickle already exist).
  • it can convert both data layouts predictions -> full_data and vice versa

I know this might be a bit of effort, but since you are refactoring it here anyway, this might pay off quickly: one place to test, and a natural spot to plug in a different storage backend later. Curious how this fits with your plans.

Comment thread deeplabcut/utils/auxfun_multianimal.py Resolved
@Ayskin
Ayskin force-pushed the yann/fix-shelving-load branch from cab117b to c902f81 Compare October 27, 2026 08:59
@Ayskin
Ayskin marked this pull request as ready for review October 28, 2026 16:13
@Ayskin
Ayskin marked this pull request as draft October 28, 2026 16:36
@Ayskin
Ayskin marked this pull request as ready for review October 29, 2026 09:29
@deruyter92
deruyter92 self-requested a review October 29, 2026 12:26

@deruyter92 deruyter92 left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Great PR, looks good. Great that you managed to centralize helpers in a context manager. Call sites look cleaner now. Just a few minor comments.

data = pickle.load(handle)
except (pickle.UnpicklingError, FileNotFoundError):
data = shelve.open(data_file, flag="r")
data = auxiliaryfunctions.load_full_data(data_file)

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Before this change, if _full.pickle wasn't a pickle but a shelve, the except branch returned a lazy object. Now the whole shelf is turned into a dict.

if not overwrite and existing_format is not None:
print(f"Video {video} already analyzed at {output_pkl}!")
else:
if existing_format is not None and bool(existing_format) != bool(use_shelve):

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Small edge case: with overwrite=True, the existing full data is deleted before we know whether the video can be read. Before, this path left previous results untouched.

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Not sure if worth addressing in this PR. Only if you find an easy solution

_DBM_SUFFIXES = (".db", ".pag", ".dir", ".dat", ".bak")


@contextmanager

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

What do you think of splitting open_full_data into a plain opener + context manager? Then we can still use the shelf object if needed:

def read_full_data(path: str | Path) -> Mapping:
    """Returns a dict for pickles; shelves are opened read-only and lazily (caller closes)."""
    path = Path(path)
    kind = full_data_format(path)
    if kind is None:
        raise FileNotFoundError(f"No full data found at {path}")
    if kind == "":
        return read_pickle(path)
    return shelve.open(str(path), flag="r")


@contextmanager
def open_full_data(path: str | Path) -> Generator[Mapping]:
    data = read_full_data(path)
    try:
        yield data
    finally:
        if isinstance(data, shelve.Shelf):
            data.close()

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants