SLEAP: Get a summary CSV of all slp files in a directory

sleap_slp_summary_csv.py

"""Outline number of labeled frames, recording sessions, and video paths for all SLPs in directory.

This script recursively finds SLP files in a directory and outputs a CSV outlining how many labeled frames, recording
sessions, and videos there are in each slp project within a directory.
"""


import pathlib

import pandas as pd

import sleap

directory = r"path/to/directory/containing/many/project.slp"

# Recursively find all SLP files in directory
slp_files = list(pathlib.Path(directory).rglob("*.slp"))
print(f"Found {len(slp_files)} SLP files in {directory}")

# Keep track of which SLP files reference same video
video_to_slp = {}

# Load each SLP file and get number of labeled frames, recording sessions, and video paths
slp_info = {}
for slp in slp_files:
    labels = sleap.load_file(slp.as_posix(), search_paths=r"path/to/find/videos/if/a/different/drive/mapping")

    n_user_annotations = len(labels.user_labeled_frames)
    recording_sessions = len(labels.sessions)

    for video in labels.videos:
        video_filename = video.filename

        if video_filename not in video_to_slp:
            video_to_slp[video_filename] = []

        # Record which SLP files reference this
        video_to_slp[video_filename].append(slp)

    # Record info
    slp_info[slp] = (n_user_annotations, recording_sessions)
    print(f"{slp}: {n_user_annotations} ground truth, {recording_sessions} recording sessions")

# Output to CSV
df_slps = pd.DataFrame(slp_info, index=["n_user_annotations", "recording_sessions"]).T
df_slps.to_csv("slp_info.csv")

df_videos = pd.DataFrame(video_to_slp).T
df_videos.to_csv("video_to_slp.csv")

print(f"Video to slp: {video_to_slp}")
print()