Tool-using mice - Lockbox#

Dataset from Reiske et al., 2025¹, containing video and pose files of individual mice solving mechanical puzzle “lockboxes” recorded from three camera perspectives (top, front, side).

¹ Reiske, P., Boon, M. N., Andresen, N., Traverso, S., Hohlbaum, K., Lewejohann, L., Thöne-Reineke, C., Hellwich, O., & Sprekeler, H. (2025). Mouse Lockbox Dataset: Behavior Recognition for Mice Solving Lockboxes (arXiv:2505.15408). arXiv. https://doi.org/10.48550/arXiv.2505.15408

../_images/lockbox1.png

From Fig. 1 in Reiske et al., 2025¹.

import requests
import zipfile
import pandas as pd
import xarray as xr
from pathlib import Path
from movement.kinematics import compute_velocity, compute_speed
from movement.utils.vector import compute_norm
from movement.io import load_poses, save_poses

import ethograph as eto
from ethograph.io.nwb_alignment import align_media_per_trial

Download data#

def download_and_extract(url: str, data_folder: Path) -> None:
    zip_path = data_folder / "dataset.zip"
    data_folder.mkdir(parents=True, exist_ok=True)
    response = requests.get(url, stream=True)
    response.raise_for_status()
    with open(zip_path, "wb") as f:
        for chunk in response.iter_content(chunk_size=8192):
            f.write(chunk)
    with zipfile.ZipFile(zip_path, "r") as zip_ref:
        zip_ref.extractall(data_folder)
    zip_path.unlink()


try:
    _here = Path(__vsc_ipynb_file__).parent
except NameError:
    _here = Path().resolve()

data_folder = _here.parent / "data" / "lockbox"
url = "https://www.dropbox.com/scl/fo/h7nkai8574h23qfq9m1b2/AP4gNZOpDJJ7z0yGtbWQiOc?rlkey=w36jzxqjkghg0j0xva5zsxy2v&e=1&st=5r9msqjw&dl=1"
download_and_extract(url, data_folder)

data_folder = data_folder / "labeled"

Build NWB alignment#

fps = 30
cameras = ["front-view", "side-view", "top-down-view"]

trials = [
    "2021-02-15_07-32-44_segment1",
    "2021-03-05_08-36-42_segment1",
    "2021-05-24_07-36-05_segment1",
    "2021-05-25_08-19-50_segment2",
    "2021-05-31_07-34-21_segment2",
    "2021-05-31_07-34-21_segment3",
]

# Discover files for each trial and camera
video_by_trial: dict[str, dict[str, str]] = {}
pose_by_trial: dict[str, dict[str, str]] = {}

for trial in trials:
    for cam in cameras:
        # Video: e.g. 2021-02-15_07-32-44_segment1_front-view.avi
        video_file = data_folder / f"{trial}_{cam}.avi"
        if video_file.exists():
            video_by_trial.setdefault(trial, {})[cam] = str(video_file)

        # Pose: convert .h5 → .csv first (done during feature extraction below),
        # stored as *_individual_0.csv
        pose_csv = data_folder / f"{trial}_{cam}_individual_0.csv"
        if pose_csv.exists():
            pose_by_trial.setdefault(trial, {})[cam] = str(pose_csv)

# Build session table
session_table = pd.DataFrame({"trial": trials})
for cam in cameras:
    session_table[f"video_{cam}"] = [
        video_by_trial.get(t, {}).get(cam, "") for t in trials
    ]
    session_table[f"pose_{cam}"] = [
        pose_by_trial.get(t, {}).get(cam, "") for t in trials
    ]
session_table = session_table.loc[:, (session_table != "").any()]

nwb_path = data_folder / ".ethograph" / "alignment.nwb"
align_media_per_trial(
    trial_table=session_table,
    stream_rates={"video": float(fps), "pose": float(fps)},
    output_path=nwb_path,
    pose_fps=float(fps),
)

Build feature datasets#

ds_list = []

for trial in trials:
    files = list(data_folder.glob(f"{trial}*"))
    trial_datasets: dict[str, xr.Dataset] = {}

    for file in files:
        name = file.name

        if file.suffix != ".h5":
            continue

        df = pd.read_hdf(file)
        ds = load_poses.from_dlc_style_df(df, fps=fps)

        # Save CSV for future use
        csv_path = str(file).replace(".h5", ".csv")
        save_poses.to_dlc_file(ds, csv_path)

        if "front-view" in name:
            ds["front_velocity"] = compute_velocity(ds.position)
            ds["front_speed"] = compute_speed(ds.position)
            trial_datasets["front-view"] = ds

        elif "top-down-view" in name:
            head_centre_pos = ds.position.sel(
                keypoints=["ear_left", "ear_right"]
            ).mean("keypoints")
            ds["topview_distance_head_lever_tip"] = compute_norm(
                ds.position.sel(keypoints="lever_tip") - head_centre_pos
            )
            ds["topview_distance_head_stick_head"] = compute_norm(
                ds.position.sel(keypoints="stick_head") - head_centre_pos
            )
            ds["topview_distance_head_ball"] = compute_norm(
                ds.position.sel(keypoints="ball") - head_centre_pos
            )
            trial_datasets["top-down-view"] = ds

        elif "side-view" in name:
            trial_datasets["side-view"] = ds

    if not trial_datasets:
        continue

    ds_merged = xr.merge(trial_datasets.values(), compat="override")
    if "top-down-view" in trial_datasets:
        ds_merged["position"] = trial_datasets["top-down-view"]["position"]
    ds_merged.attrs["trial"] = trial
    ds_list.append(ds_merged)

dt = eto.from_datasets(ds_list)
dt.save(data_folder / "lockbox.nc")
print(f"Saved to {data_folder / 'lockbox.nc'}")