Skip to content

Refactor/predict multianimal #3220

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
MMathisLab merged 7 commits into from
May 6, 2026
Merged

Refactor/predict multianimal #3220

Changes from all commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
8ca148f
refactor: modernize path handling in predict_multianimal.py
juan-cobos Feb 25, 2026
a237689
Remove unnecessary auxfunc import
juan-cobos Feb 25, 2026
30b513f
Fix .h5 saving in auxfun_multianimal.SaveFullMultiAnimalData, write p…
juan-cobos Feb 27, 2026
0360b3b
Fix typo of results dict
juan-cobos Feb 27, 2026
e65210c
Close context manager properly and inverted batch_size <= 1 error logic
juan-cobos Feb 27, 2026
ad474c3
Run pre-commit
C-Achard Mar 30, 2026
99e810a
Merge branch 'main' into refactor/predict-multianimal
C-Achard Apr 8, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
308 changes: 148 additions & 160 deletions deeplabcut/pose_estimation_tensorflow/predict_multianimal.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +9,6 @@
# Licensed under GNU Lesser General Public License v3.0
#

import os
import pickle
import shelve
import time
Expand All @@ -21,7 +20,7 @@
from tqdm import tqdm

from deeplabcut.pose_estimation_tensorflow.core import predict_multianimal as predict
from deeplabcut.utils import auxfun_multianimal, auxiliaryfunctions
from deeplabcut.utils import auxfun_multianimal
Comment thread
juan-cobos marked this conversation as resolved.
from deeplabcut.utils.auxfun_videos import VideoWriter


Expand All @@ -39,68 +38,61 @@ def extract_bpt_feature_from_video(
robust_nframes=False,
):
print("Starting to analyze % ", video)
vname = Path(video).stem
videofolder = str(Path(video).parents[0])
if destfolder is None:
destfolder = videofolder
auxiliaryfunctions.attempt_to_make_folder(destfolder)
dataname = os.path.join(destfolder, vname + DLCscorer + ".h5")

assemble_filename = dataname.split(".h5")[0] + "_assemblies.pickle"
video = Path(video)
destfolder = video.parent if destfolder is None else Path(destfolder)
destfolder.mkdir(exist_ok=True, parents=True)
basename = f"{video.stem}{DLCscorer}"

feature_dict = shelve.open(
dataname.split(".h5")[0] + "_bpt_features.pickle",
str(destfolder / f"{basename}_bpt_features.pickle"),
protocol=pickle.DEFAULT_PROTOCOL,
)

with open(assemble_filename, "rb") as f:
with open(destfolder / f"{basename}_assemblies.pickle", "rb") as f:
Comment thread
juan-cobos marked this conversation as resolved.
assemblies = pickle.load(f)
print("Loading ", video)
vid = VideoWriter(video)
if robust_nframes:
nframes = vid.get_n_frames(robust=True)
duration = vid.calc_duration(robust=True)
fps = nframes / duration
else:
nframes = len(vid)
duration = vid.calc_duration(robust=False)
fps = vid.fps

nx, ny = vid.dimensions
print(
"Duration of video [s]: ",
round(duration, 2),
", recorded with ",
round(fps, 2),
"fps!",
)
print(
"Overall # of frames: ",
nframes,
" found with (before cropping) frame dimensions: ",
nx,
ny,
)
time.time()

print("Starting to extract posture")
if int(dlc_cfg["batch_size"]) > 1:
# for multi animal, seems only this is used
PredicteData, nframes = GetPoseandCostsF_from_assemblies(
cfg,
dlc_cfg,
sess,
inputs,
outputs,
vid,
nframes,
int(dlc_cfg["batch_size"]),
assemblies,
feature_dict,
extra_dict,
)
else:
raise NotImplementedError("Not implemented yet, please raise an GitHub issue if you need this.")
print("Loading ", video)
vid = VideoWriter(str(video))
if robust_nframes:
nframes = vid.get_n_frames(robust=True)
duration = vid.calc_duration(robust=True)
fps = nframes / duration
else:
nframes = len(vid)
duration = vid.calc_duration(robust=False)
fps = vid.fps

print(
"Duration of video [s]: ",
round(duration, 2),
", recorded with ",
round(fps, 2),
"fps!",
)
print(
"Overall # of frames: ",
nframes,
" found with (before cropping) frame dimensions: ",
vid.dimensions,
)

print("Starting to extract posture")
if int(dlc_cfg["batch_size"]) <= 1:
raise NotImplementedError("Not implemented yet, please raise an GitHub issue if you need this.")
# for multi animal, seems only 'dlc_cfg["batch_size"]) > 1' is used
predicted_data, nframes = GetPoseandCostsF_from_assemblies(
cfg,
dlc_cfg,
sess,
inputs,
outputs,
vid,
nframes,
int(dlc_cfg["batch_size"]),
assemblies,
feature_dict,
extra_dict,
)


def AnalyzeMultiAnimalVideo(
Expand All @@ -116,110 +108,106 @@ def AnalyzeMultiAnimalVideo(
robust_nframes=False,
use_shelve=False,
):
"""Helper function for analyzing a video with multiple individuals."""
"""Helper function for analyzing a video with multiple individuals"""

print("Starting to analyze % ", video)
vname = Path(video).stem
videofolder = str(Path(video).parents[0])
if destfolder is None:
destfolder = videofolder
auxiliaryfunctions.attempt_to_make_folder(destfolder)
dataname = os.path.join(destfolder, vname + DLCscorer + ".h5")

if os.path.isfile(dataname.split(".h5")[0] + "_full.pickle"):
print("Video already analyzed!", dataname)
video = Path(video)
destfolder = video.parent if destfolder is None else Path(destfolder)
destfolder.mkdir(exist_ok=True, parents=True)
basename = f"{video.stem}{DLCscorer}"
full_pickle = destfolder / f"{basename}_full.pickle"

if full_pickle.is_file():
print("Video already analyzed!", full_pickle)
Comment thread
juan-cobos marked this conversation as resolved.
Comment thread
juan-cobos marked this conversation as resolved.
return None

print("Loading ", video)
vid = VideoWriter(str(video))
if robust_nframes:
nframes = vid.get_n_frames(robust=True)
duration = vid.calc_duration(robust=True)
fps = nframes / duration
else:
print("Loading ", video)
vid = VideoWriter(video)
if robust_nframes:
nframes = vid.get_n_frames(robust=True)
duration = vid.calc_duration(robust=True)
fps = nframes / duration
else:
nframes = len(vid)
duration = vid.calc_duration(robust=False)
fps = vid.fps

nx, ny = vid.dimensions
print(
"Duration of video [s]: ",
round(duration, 2),
", recorded with ",
round(fps, 2),
"fps!",
nframes = len(vid)
duration = vid.calc_duration(robust=False)
fps = vid.fps

print(
"Duration of video [s]: ",
round(duration, 2),
", recorded with ",
round(fps, 2),
"fps!",
)
print(
"Overall # of frames: ",
nframes,
" found with (before cropping) frame dimensions: ",
vid.dimensions,
)
start = time.time()

print(
"Starting to extract posture from the video(s) with batchsize:",
dlc_cfg["batch_size"],
)

shelf_path = str(full_pickle) if use_shelve else ""
if int(dlc_cfg["batch_size"]) > 1:
predicted_data, nframes = GetPoseandCostsF(
cfg,
dlc_cfg,
sess,
inputs,
outputs,
vid,
nframes,
int(dlc_cfg["batch_size"]),
shelf_path,
)
print(
"Overall # of frames: ",
else:
predicted_data, nframes = GetPoseandCostsS(
cfg,
dlc_cfg,
sess,
inputs,
outputs,
vid,
nframes,
" found with (before cropping) frame dimensions: ",
nx,
ny,
shelf_path,
)
start = time.time()

print(
"Starting to extract posture from the video(s) with batchsize:",
dlc_cfg["batch_size"],
)
if use_shelve:
shelf_path = dataname.split(".h5")[0] + "_full.pickle"
else:
shelf_path = ""
if int(dlc_cfg["batch_size"]) > 1:
PredicteData, nframes = GetPoseandCostsF(
cfg,
dlc_cfg,
sess,
inputs,
outputs,
vid,
nframes,
int(dlc_cfg["batch_size"]),
shelf_path,
)
else:
PredicteData, nframes = GetPoseandCostsS(
cfg,
dlc_cfg,
sess,
inputs,
outputs,
vid,
nframes,
shelf_path,
)
stop = time.time()

nx, ny = vid.dimensions
if cfg["cropping"]:
Comment thread
juan-cobos marked this conversation as resolved.
coords = [cfg["x1"], cfg["x2"], cfg["y1"], cfg["y2"]]
else:
coords = [0, nx, 0, ny]

dictionary = {
"start": start,
"stop": stop,
"run_duration": stop - start,
"Scorer": DLCscorer,
"DLC-model-config file": dlc_cfg,
"fps": fps,
"batch_size": dlc_cfg["batch_size"],
"frame_dimensions": (ny, nx),
"nframes": nframes,
"iteration (active-learning)": cfg["iteration"],
"training set fraction": trainFraction,
"cropping": cfg["cropping"],
"cropping_parameters": coords,
}
metadata = {"data": dictionary}
print(f"Video Analyzed. Saving results in {destfolder}")

stop = time.time()

if cfg["cropping"]:
coords = [cfg["x1"], cfg["x2"], cfg["y1"], cfg["y2"]]
else:
coords = [0, nx, 0, ny]

dictionary = {
"start": start,
"stop": stop,
"run_duration": stop - start,
"Scorer": DLCscorer,
"DLC-model-config file": dlc_cfg,
"fps": fps,
"batch_size": dlc_cfg["batch_size"],
"frame_dimensions": (ny, nx),
"nframes": nframes,
"iteration (active-learning)": cfg["iteration"],
"training set fraction": trainFraction,
"cropping": cfg["cropping"],
"cropping_parameters": coords,
}
metadata = {"data": dictionary}
print(f"Video Analyzed. Saving results in {destfolder}...")

if use_shelve:
metadata_path = dataname.split(".h5")[0] + "_meta.pickle"
with open(metadata_path, "wb") as f:
pickle.dump(metadata, f, pickle.HIGHEST_PROTOCOL)
else:
_ = auxfun_multianimal.SaveFullMultiAnimalData(PredicteData, metadata, dataname)
if use_shelve:
with open(destfolder / f"{basename}_meta.pickle", "wb") as f:
pickle.dump(metadata, f, pickle.HIGHEST_PROTOCOL)
else:
auxfun_multianimal.SaveFullMultiAnimalData(predicted_data, metadata, str(destfolder / f"{basename}.h5"))


def _get_features_dict(raw_coords, features, stride):
Expand Down Expand Up @@ -252,7 +240,7 @@ def GetPoseandCostsF_from_assemblies(
feature_dict,
extra_dict,
):
"""Batchwise prediction of pose."""
"""Batchwise prediction of pose"""
strwidth = int(np.ceil(np.log10(nframes))) # width for strings
batch_ind = 0 # keeps track of which image within a batch should be written to
batch_num = 0 # keeps track of which batch you are at
Expand All @@ -265,7 +253,7 @@ def GetPoseandCostsF_from_assemblies(
counter = 0
inds = []

PredicteData = {}
predicted_data = {}

while cap.video.isOpened():
frame = cap.read_frame(crop=cfg["cropping"])
Expand All @@ -290,7 +278,7 @@ def GetPoseandCostsF_from_assemblies(

D, features = preds
for i, (ind, data) in enumerate(zip(inds, D, strict=False)):
PredicteData["frame" + str(ind).zfill(strwidth)] = data
predicted_data["frame" + str(ind).zfill(strwidth)] = data
raw_coords = assemblies.get(ind)
if raw_coords is None:
continue
Expand All @@ -316,7 +304,7 @@ def GetPoseandCostsF_from_assemblies(

D, features = preds
for i, (ind, data) in enumerate(zip(inds, D, strict=False)):
PredicteData["frame" + str(ind).zfill(strwidth)] = data
predicted_data["frame" + str(ind).zfill(strwidth)] = data
raw_coords = assemblies.get(ind)
if raw_coords is None:
continue
Expand All @@ -334,7 +322,7 @@ def GetPoseandCostsF_from_assemblies(
cap.close()
pbar.close()
feature_dict.close()
PredicteData["metadata"] = {
predicted_data["metadata"] = {
"nms radius": dlc_cfg["nmsradius"],
"minimal confidence": dlc_cfg["minconfidence"],
"sigma": dlc_cfg.get("sigma", 1),
Expand All @@ -344,7 +332,7 @@ def GetPoseandCostsF_from_assemblies(
"all_joints_names": [dlc_cfg["all_joints_names"][i] for i in range(len(dlc_cfg["all_joints"]))],
"nframes": nframes,
}
return PredicteData, nframes
return predicted_data, nframes


def GetPoseandCostsF(
Expand All @@ -358,7 +346,7 @@ def GetPoseandCostsF(
batchsize,
shelf_path,
):
"""Batchwise prediction of pose."""
"""Batchwise prediction of pose"""
strwidth = int(np.ceil(np.log10(nframes))) # width for strings
batch_ind = 0 # keeps track of which image within a batch should be written to
batch_num = 0 # keeps track of which batch you are at
Expand Down
Loading