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

Open
juan-cobos wants to merge 7 commits into
base: main
Choose a base branch
from
+150 −163
Open

Refactor/predict multianimal #3220

Changes from all commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
a87945e
refactor: modernize path handling in predict_multianimal.py
juan-cobos Feb 25, 2026
29dbf0e
Merge branch 'DeepLabCut:main' into refactor/predict-multianimal
juan-cobos Feb 25, 2026
c61d853
Remove unnecessary auxfunc import
juan-cobos Feb 25, 2026
46410ff
Fix .h5 saving in auxfun_multianimal.SaveFullMultiAnimalData, write p…
juan-cobos Feb 27, 2026
2da1af4
Fix typo of results dict
juan-cobos Feb 27, 2026
0a5a3fd
Close context manager properly and inverted batch_size <= 1 error logic
juan-cobos Feb 27, 2026
cc25c74
Merge branch 'main' into refactor/predict-multianimal
juan-cobos Mar 12, 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
313 changes: 150 additions & 163 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,9 +20,8 @@
from tqdm import tqdm

from deeplabcut.pose_estimation_tensorflow.core import predict_multianimal as predict
from deeplabcut.utils import auxiliaryfunctions, auxfun_multianimal
from deeplabcut.utils import auxfun_multianimal
from deeplabcut.utils.auxfun_videos import VideoWriter
import pickle


def extract_bpt_feature_from_video(
Expand All @@ -40,70 +38,63 @@ 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:
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,
)
start = 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 @@ -122,115 +113,111 @@ def AnalyzeMultiAnimalVideo(
"""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)
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()

stop = time.time()

if cfg["cropping"] == True:
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("Video Analyzed. Saving results in %s..." % (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
)
nx, ny = vid.dimensions
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:
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):
from deeplabcut.pose_tracking_pytorch import (
load_features_from_coord,
convert_coord_from_img_space_to_feature_space,
load_features_from_coord,
)

coords_img_space = np.array(
Expand Down Expand Up @@ -276,7 +263,7 @@ def GetPoseandCostsF_from_assemblies(
counter = 0
inds = []

PredicteData = {}
predicted_data = {}

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

D, features = preds
for i, (ind, data) in enumerate(zip(inds, D)):
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 @@ -327,7 +314,7 @@ def GetPoseandCostsF_from_assemblies(

D, features = preds
for i, (ind, data) in enumerate(zip(inds, D)):
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 @@ -345,7 +332,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 @@ -359,7 +346,7 @@ def GetPoseandCostsF_from_assemblies(
],
"nframes": nframes,
}
return PredicteData, nframes
return predicted_data, nframes


def GetPoseandCostsF(
Expand Down
Loading