From 8ca148f530255b740e5d3b68c40ebe7282c9dd9c Mon Sep 17 00:00:00 2001 From: Juan Cobos Date: Wed, 25 Feb 2026 13:32:52 +0100 Subject: [PATCH 1/6] refactor: modernize path handling in predict_multianimal.py - Replace os.path string manipulation with pathlib throughout - Drop duplicate import pickle and unused import os - Guard destfolder as Path in both functions - Use VideoWriter(str(video)) to avoid cv2.VideoCapture breakage - Extract full_pickle variable to avoid repeating the _full.pickle path - Early return in AnalyzeMultiAnimalVideo instead of deep else nesting - Inline vid.dimensions in prints; keep nx, ny only where needed - shelf_path ternary; f-string print; drop unused _ assignment --- .../predict_multianimal.py | 262 +++++++++--------- 1 file changed, 136 insertions(+), 126 deletions(-) diff --git a/deeplabcut/pose_estimation_tensorflow/predict_multianimal.py b/deeplabcut/pose_estimation_tensorflow/predict_multianimal.py index a8b7248c95..e8197e75ba 100644 --- a/deeplabcut/pose_estimation_tensorflow/predict_multianimal.py +++ b/deeplabcut/pose_estimation_tensorflow/predict_multianimal.py @@ -9,7 +9,6 @@ # Licensed under GNU Lesser General Public License v3.0 # -import os import pickle import shelve import time @@ -39,24 +38,20 @@ 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) + dataname = destfolder / f"{video.stem}{DLCscorer}" feature_dict = shelve.open( - dataname.split(".h5")[0] + "_bpt_features.pickle", + f"{dataname}_bpt_features.pickle", protocol=pickle.DEFAULT_PROTOCOL, ) - with open(assemble_filename, "rb") as f: + with open(f"{dataname}_assemblies.pickle", "rb") as f: assemblies = pickle.load(f) print("Loading ", video) - vid = VideoWriter(video) + vid = VideoWriter(str(video)) if robust_nframes: nframes = vid.get_n_frames(robust=True) duration = vid.calc_duration(robust=True) @@ -66,7 +61,6 @@ def extract_bpt_feature_from_video( duration = vid.calc_duration(robust=False) fps = vid.fps - nx, ny = vid.dimensions print( "Duration of video [s]: ", round(duration, 2), @@ -78,10 +72,8 @@ def extract_bpt_feature_from_video( "Overall # of frames: ", nframes, " found with (before cropping) frame dimensions: ", - nx, - ny, + vid.dimensions, ) - time.time() print("Starting to extract posture") if int(dlc_cfg["batch_size"]) > 1: @@ -100,7 +92,9 @@ def extract_bpt_feature_from_video( extra_dict, ) else: - raise NotImplementedError("Not implemented yet, please raise an GitHub issue if you need this.") + raise NotImplementedError( + "Not implemented yet, please raise an GitHub issue if you need this." + ) def AnalyzeMultiAnimalVideo( @@ -116,110 +110,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"): + video = Path(video) + destfolder = video.parent if destfolder is None else Path(destfolder) + destfolder.mkdir(exist_ok=True, parents=True) + dataname = destfolder / f"{video.stem}{DLCscorer}" + full_pickle = Path(f"{dataname}_full.pickle") + + if full_pickle.is_file(): print("Video already analyzed!", dataname) + 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 + 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() - nx, ny = vid.dimensions - print( - "Duration of video [s]: ", - round(duration, 2), - ", recorded with ", - round(fps, 2), - "fps!", + 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: + PredicteData, nframes = GetPoseandCostsF( + cfg, + dlc_cfg, + sess, + inputs, + outputs, + vid, + nframes, + int(dlc_cfg["batch_size"]), + shelf_path, ) - print( - "Overall # of frames: ", + else: + PredicteData, 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() + 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 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(f"{dataname}_meta.pickle", "wb") as f: + pickle.dump(metadata, f, pickle.HIGHEST_PROTOCOL) + else: + auxfun_multianimal.SaveFullMultiAnimalData(PredicteData, metadata, str(dataname)) def _get_features_dict(raw_coords, features, stride): @@ -228,14 +218,18 @@ def _get_features_dict(raw_coords, features, stride): load_features_from_coord, ) - coords_img_space = np.array([coord[:, :2] for coord in raw_coords]) # only first two columns are useful + coords_img_space = np.array( + [coord[:, :2] for coord in raw_coords] + ) # only first two columns are useful coords_feature_space = convert_coord_from_img_space_to_feature_space( coords_img_space, stride, ) - bpt_features = load_features_from_coord(features.astype(np.float16), coords_feature_space) + bpt_features = load_features_from_coord( + features.astype(np.float16), coords_feature_space + ) return {"features": bpt_features, "coordinates": coords_img_space} @@ -252,7 +246,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 @@ -260,7 +254,9 @@ def GetPoseandCostsF_from_assemblies( cap.set_bbox(cfg["x1"], cfg["x2"], cfg["y1"], cfg["y2"]) nx, ny = cap.dimensions - frames = np.empty((batchsize, ny, nx, 3), dtype="ubyte") # this keeps all frames in a batch + frames = np.empty( + (batchsize, ny, nx, 3), dtype="ubyte" + ) # this keeps all frames in a batch pbar = tqdm(total=nframes) counter = 0 inds = [] @@ -289,7 +285,7 @@ def GetPoseandCostsF_from_assemblies( continue D, features = preds - for i, (ind, data) in enumerate(zip(inds, D, strict=False)): + for i, (ind, data) in enumerate(zip(inds, D)): PredicteData["frame" + str(ind).zfill(strwidth)] = data raw_coords = assemblies.get(ind) if raw_coords is None: @@ -315,7 +311,7 @@ def GetPoseandCostsF_from_assemblies( continue D, features = preds - for i, (ind, data) in enumerate(zip(inds, D, strict=False)): + for i, (ind, data) in enumerate(zip(inds, D)): PredicteData["frame" + str(ind).zfill(strwidth)] = data raw_coords = assemblies.get(ind) if raw_coords is None: @@ -339,9 +335,13 @@ def GetPoseandCostsF_from_assemblies( "minimal confidence": dlc_cfg["minconfidence"], "sigma": dlc_cfg.get("sigma", 1), "PAFgraph": dlc_cfg["partaffinityfield_graph"], - "PAFinds": dlc_cfg.get("paf_best", np.arange(len(dlc_cfg["partaffinityfield_graph"]))), + "PAFinds": dlc_cfg.get( + "paf_best", np.arange(len(dlc_cfg["partaffinityfield_graph"])) + ), "all_joints": [[i] for i in range(len(dlc_cfg["all_joints"]))], - "all_joints_names": [dlc_cfg["all_joints_names"][i] for i in range(len(dlc_cfg["all_joints"]))], + "all_joints_names": [ + dlc_cfg["all_joints_names"][i] for i in range(len(dlc_cfg["all_joints"])) + ], "nframes": nframes, } return PredicteData, nframes @@ -358,7 +358,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 @@ -366,7 +366,9 @@ def GetPoseandCostsF( cap.set_bbox(cfg["x1"], cfg["x2"], cfg["y1"], cfg["y2"]) nx, ny = cap.dimensions - frames = np.empty((batchsize, ny, nx, 3), dtype="ubyte") # this keeps all frames in a batch + frames = np.empty( + (batchsize, ny, nx, 3), dtype="ubyte" + ) # this keeps all frames in a batch pbar = tqdm(total=nframes) counter = 0 inds = [] @@ -383,9 +385,13 @@ def GetPoseandCostsF( "minimal confidence": dlc_cfg["minconfidence"], "sigma": dlc_cfg.get("sigma", 1), "PAFgraph": dlc_cfg["partaffinityfield_graph"], - "PAFinds": dlc_cfg.get("paf_best", np.arange(len(dlc_cfg["partaffinityfield_graph"]))), + "PAFinds": dlc_cfg.get( + "paf_best", np.arange(len(dlc_cfg["partaffinityfield_graph"])) + ), "all_joints": [[i] for i in range(len(dlc_cfg["all_joints"]))], - "all_joints_names": [dlc_cfg["all_joints_names"][i] for i in range(len(dlc_cfg["all_joints"]))], + "all_joints_names": [ + dlc_cfg["all_joints_names"][i] for i in range(len(dlc_cfg["all_joints"])) + ], "nframes": nframes, } while cap.video.isOpened(): @@ -408,7 +414,7 @@ def GetPoseandCostsF( inputs, outputs, ) - for ind, data in zip(inds, D, strict=False): + for ind, data in zip(inds, D): db["frame" + str(ind).zfill(strwidth)] = data del D batch_ind = 0 @@ -425,7 +431,7 @@ def GetPoseandCostsF( inputs, outputs, ) - for ind, data in zip(inds, D, strict=False): + for ind, data in zip(inds, D): db["frame" + str(ind).zfill(strwidth)] = data del D break @@ -459,9 +465,13 @@ def GetPoseandCostsS(cfg, dlc_cfg, sess, inputs, outputs, cap, nframes, shelf_pa "minimal confidence": dlc_cfg["minconfidence"], "sigma": dlc_cfg.get("sigma", 1), "PAFgraph": dlc_cfg["partaffinityfield_graph"], - "PAFinds": dlc_cfg.get("paf_best", np.arange(len(dlc_cfg["partaffinityfield_graph"]))), + "PAFinds": dlc_cfg.get( + "paf_best", np.arange(len(dlc_cfg["partaffinityfield_graph"])) + ), "all_joints": [[i] for i in range(len(dlc_cfg["all_joints"]))], - "all_joints_names": [dlc_cfg["all_joints_names"][i] for i in range(len(dlc_cfg["all_joints"]))], + "all_joints_names": [ + dlc_cfg["all_joints_names"][i] for i in range(len(dlc_cfg["all_joints"])) + ], "nframes": nframes, } pbar = tqdm(total=nframes) From a23768901d07fd4f3cdb4662c4d1bb179b10c4a0 Mon Sep 17 00:00:00 2001 From: Juan Cobos Date: Wed, 25 Feb 2026 16:34:24 +0100 Subject: [PATCH 2/6] Remove unnecessary auxfunc import --- .../pose_estimation_tensorflow/predict_multianimal.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/deeplabcut/pose_estimation_tensorflow/predict_multianimal.py b/deeplabcut/pose_estimation_tensorflow/predict_multianimal.py index e8197e75ba..af2ba7f77b 100644 --- a/deeplabcut/pose_estimation_tensorflow/predict_multianimal.py +++ b/deeplabcut/pose_estimation_tensorflow/predict_multianimal.py @@ -20,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 from deeplabcut.utils.auxfun_videos import VideoWriter @@ -209,7 +209,9 @@ def AnalyzeMultiAnimalVideo( with open(f"{dataname}_meta.pickle", "wb") as f: pickle.dump(metadata, f, pickle.HIGHEST_PROTOCOL) else: - auxfun_multianimal.SaveFullMultiAnimalData(PredicteData, metadata, str(dataname)) + auxfun_multianimal.SaveFullMultiAnimalData( + PredicteData, metadata, str(dataname) + ) def _get_features_dict(raw_coords, features, stride): From 30b513f940b6d678262b6d950dcf8fe94064758f Mon Sep 17 00:00:00 2001 From: Juan Cobos Date: Fri, 27 Feb 2026 10:26:15 +0100 Subject: [PATCH 3/6] Fix .h5 saving in auxfun_multianimal.SaveFullMultiAnimalData, write paths as destfolder with basename instead of splitting based on extension, fix str path compatibility with shelve --- .../predict_multianimal.py | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/deeplabcut/pose_estimation_tensorflow/predict_multianimal.py b/deeplabcut/pose_estimation_tensorflow/predict_multianimal.py index af2ba7f77b..652113f4f7 100644 --- a/deeplabcut/pose_estimation_tensorflow/predict_multianimal.py +++ b/deeplabcut/pose_estimation_tensorflow/predict_multianimal.py @@ -41,14 +41,14 @@ def extract_bpt_feature_from_video( video = Path(video) destfolder = video.parent if destfolder is None else Path(destfolder) destfolder.mkdir(exist_ok=True, parents=True) - dataname = destfolder / f"{video.stem}{DLCscorer}" + basename = f"{video.stem}{DLCscorer}" feature_dict = shelve.open( - f"{dataname}_bpt_features.pickle", + str(destfolder / f"{basename}_bpt_features.pickle"), protocol=pickle.DEFAULT_PROTOCOL, ) - with open(f"{dataname}_assemblies.pickle", "rb") as f: + with open(destfolder / f"{basename}_assemblies.pickle", "rb") as f: assemblies = pickle.load(f) print("Loading ", video) vid = VideoWriter(str(video)) @@ -116,11 +116,11 @@ def AnalyzeMultiAnimalVideo( video = Path(video) destfolder = video.parent if destfolder is None else Path(destfolder) destfolder.mkdir(exist_ok=True, parents=True) - dataname = destfolder / f"{video.stem}{DLCscorer}" - full_pickle = Path(f"{dataname}_full.pickle") + basename = f"{video.stem}{DLCscorer}" + full_pickle = destfolder / f"{basename}_full.pickle" if full_pickle.is_file(): - print("Video already analyzed!", dataname) + print("Video already analyzed!", full_pickle) return None print("Loading ", video) @@ -206,11 +206,11 @@ def AnalyzeMultiAnimalVideo( print(f"Video Analyzed. Saving results in {destfolder}") if use_shelve: - with open(f"{dataname}_meta.pickle", "wb") as f: + with open(destfolder / f"{basename}_meta.pickle", "wb") as f: pickle.dump(metadata, f, pickle.HIGHEST_PROTOCOL) else: auxfun_multianimal.SaveFullMultiAnimalData( - PredicteData, metadata, str(dataname) + PredicteData, metadata, str(destfolder / f"{basename}.h5") ) From 0360b3b099c36ba0d91503c0bca6d4f18d99ec71 Mon Sep 17 00:00:00 2001 From: Juan Cobos Date: Fri, 27 Feb 2026 11:40:29 +0100 Subject: [PATCH 4/6] Fix typo of results dict --- .../predict_multianimal.py | 18 +++++++++--------- 1 file changed, 9 insertions(+), 9 deletions(-) diff --git a/deeplabcut/pose_estimation_tensorflow/predict_multianimal.py b/deeplabcut/pose_estimation_tensorflow/predict_multianimal.py index 652113f4f7..4f1cf2f82e 100644 --- a/deeplabcut/pose_estimation_tensorflow/predict_multianimal.py +++ b/deeplabcut/pose_estimation_tensorflow/predict_multianimal.py @@ -78,7 +78,7 @@ def extract_bpt_feature_from_video( 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( + predicted_data, nframes = GetPoseandCostsF_from_assemblies( cfg, dlc_cfg, sess, @@ -156,7 +156,7 @@ def AnalyzeMultiAnimalVideo( shelf_path = str(full_pickle) if use_shelve else "" if int(dlc_cfg["batch_size"]) > 1: - PredicteData, nframes = GetPoseandCostsF( + predicted_data, nframes = GetPoseandCostsF( cfg, dlc_cfg, sess, @@ -168,7 +168,7 @@ def AnalyzeMultiAnimalVideo( shelf_path, ) else: - PredicteData, nframes = GetPoseandCostsS( + predicted_data, nframes = GetPoseandCostsS( cfg, dlc_cfg, sess, @@ -210,7 +210,7 @@ def AnalyzeMultiAnimalVideo( pickle.dump(metadata, f, pickle.HIGHEST_PROTOCOL) else: auxfun_multianimal.SaveFullMultiAnimalData( - PredicteData, metadata, str(destfolder / f"{basename}.h5") + predicted_data, metadata, str(destfolder / f"{basename}.h5") ) @@ -263,7 +263,7 @@ def GetPoseandCostsF_from_assemblies( counter = 0 inds = [] - PredicteData = {} + predicted_data = {} while cap.video.isOpened(): frame = cap.read_frame(crop=cfg["cropping"]) @@ -288,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 @@ -314,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 @@ -332,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), @@ -346,7 +346,7 @@ def GetPoseandCostsF_from_assemblies( ], "nframes": nframes, } - return PredicteData, nframes + return predicted_data, nframes def GetPoseandCostsF( From e65210c3205e96025e529a50eb402504da241869 Mon Sep 17 00:00:00 2001 From: Juan Cobos Date: Fri, 27 Feb 2026 11:50:14 +0100 Subject: [PATCH 5/6] Close context manager properly and inverted batch_size <= 1 error logic --- .../predict_multianimal.py | 88 +++++++++---------- 1 file changed, 44 insertions(+), 44 deletions(-) diff --git a/deeplabcut/pose_estimation_tensorflow/predict_multianimal.py b/deeplabcut/pose_estimation_tensorflow/predict_multianimal.py index 4f1cf2f82e..c9b9bfdae1 100644 --- a/deeplabcut/pose_estimation_tensorflow/predict_multianimal.py +++ b/deeplabcut/pose_estimation_tensorflow/predict_multianimal.py @@ -50,51 +50,51 @@ def extract_bpt_feature_from_video( with open(destfolder / f"{basename}_assemblies.pickle", "rb") as f: assemblies = pickle.load(f) - 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: - # for multi animal, seems only this 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, - ) - 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( From ad474c36e86356e72fc2cfc2b2061d0ca9f5ad85 Mon Sep 17 00:00:00 2001 From: Cyril Achard Date: Mon, 30 Mar 2026 13:41:48 -0500 Subject: [PATCH 6/6] Run pre-commit --- .../predict_multianimal.py | 56 ++++++------------- 1 file changed, 16 insertions(+), 40 deletions(-) diff --git a/deeplabcut/pose_estimation_tensorflow/predict_multianimal.py b/deeplabcut/pose_estimation_tensorflow/predict_multianimal.py index c9b9bfdae1..00c23b24dd 100644 --- a/deeplabcut/pose_estimation_tensorflow/predict_multianimal.py +++ b/deeplabcut/pose_estimation_tensorflow/predict_multianimal.py @@ -78,9 +78,7 @@ def extract_bpt_feature_from_video( 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." - ) + 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, @@ -209,9 +207,7 @@ def AnalyzeMultiAnimalVideo( 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") - ) + auxfun_multianimal.SaveFullMultiAnimalData(predicted_data, metadata, str(destfolder / f"{basename}.h5")) def _get_features_dict(raw_coords, features, stride): @@ -220,18 +216,14 @@ def _get_features_dict(raw_coords, features, stride): load_features_from_coord, ) - coords_img_space = np.array( - [coord[:, :2] for coord in raw_coords] - ) # only first two columns are useful + coords_img_space = np.array([coord[:, :2] for coord in raw_coords]) # only first two columns are useful coords_feature_space = convert_coord_from_img_space_to_feature_space( coords_img_space, stride, ) - bpt_features = load_features_from_coord( - features.astype(np.float16), coords_feature_space - ) + bpt_features = load_features_from_coord(features.astype(np.float16), coords_feature_space) return {"features": bpt_features, "coordinates": coords_img_space} @@ -256,9 +248,7 @@ def GetPoseandCostsF_from_assemblies( cap.set_bbox(cfg["x1"], cfg["x2"], cfg["y1"], cfg["y2"]) nx, ny = cap.dimensions - frames = np.empty( - (batchsize, ny, nx, 3), dtype="ubyte" - ) # this keeps all frames in a batch + frames = np.empty((batchsize, ny, nx, 3), dtype="ubyte") # this keeps all frames in a batch pbar = tqdm(total=nframes) counter = 0 inds = [] @@ -287,7 +277,7 @@ def GetPoseandCostsF_from_assemblies( continue D, features = preds - for i, (ind, data) in enumerate(zip(inds, D)): + for i, (ind, data) in enumerate(zip(inds, D, strict=False)): predicted_data["frame" + str(ind).zfill(strwidth)] = data raw_coords = assemblies.get(ind) if raw_coords is None: @@ -313,7 +303,7 @@ def GetPoseandCostsF_from_assemblies( continue D, features = preds - for i, (ind, data) in enumerate(zip(inds, D)): + for i, (ind, data) in enumerate(zip(inds, D, strict=False)): predicted_data["frame" + str(ind).zfill(strwidth)] = data raw_coords = assemblies.get(ind) if raw_coords is None: @@ -337,13 +327,9 @@ def GetPoseandCostsF_from_assemblies( "minimal confidence": dlc_cfg["minconfidence"], "sigma": dlc_cfg.get("sigma", 1), "PAFgraph": dlc_cfg["partaffinityfield_graph"], - "PAFinds": dlc_cfg.get( - "paf_best", np.arange(len(dlc_cfg["partaffinityfield_graph"])) - ), + "PAFinds": dlc_cfg.get("paf_best", np.arange(len(dlc_cfg["partaffinityfield_graph"]))), "all_joints": [[i] for i in range(len(dlc_cfg["all_joints"]))], - "all_joints_names": [ - dlc_cfg["all_joints_names"][i] for i in range(len(dlc_cfg["all_joints"])) - ], + "all_joints_names": [dlc_cfg["all_joints_names"][i] for i in range(len(dlc_cfg["all_joints"]))], "nframes": nframes, } return predicted_data, nframes @@ -368,9 +354,7 @@ def GetPoseandCostsF( cap.set_bbox(cfg["x1"], cfg["x2"], cfg["y1"], cfg["y2"]) nx, ny = cap.dimensions - frames = np.empty( - (batchsize, ny, nx, 3), dtype="ubyte" - ) # this keeps all frames in a batch + frames = np.empty((batchsize, ny, nx, 3), dtype="ubyte") # this keeps all frames in a batch pbar = tqdm(total=nframes) counter = 0 inds = [] @@ -387,13 +371,9 @@ def GetPoseandCostsF( "minimal confidence": dlc_cfg["minconfidence"], "sigma": dlc_cfg.get("sigma", 1), "PAFgraph": dlc_cfg["partaffinityfield_graph"], - "PAFinds": dlc_cfg.get( - "paf_best", np.arange(len(dlc_cfg["partaffinityfield_graph"])) - ), + "PAFinds": dlc_cfg.get("paf_best", np.arange(len(dlc_cfg["partaffinityfield_graph"]))), "all_joints": [[i] for i in range(len(dlc_cfg["all_joints"]))], - "all_joints_names": [ - dlc_cfg["all_joints_names"][i] for i in range(len(dlc_cfg["all_joints"])) - ], + "all_joints_names": [dlc_cfg["all_joints_names"][i] for i in range(len(dlc_cfg["all_joints"]))], "nframes": nframes, } while cap.video.isOpened(): @@ -416,7 +396,7 @@ def GetPoseandCostsF( inputs, outputs, ) - for ind, data in zip(inds, D): + for ind, data in zip(inds, D, strict=False): db["frame" + str(ind).zfill(strwidth)] = data del D batch_ind = 0 @@ -433,7 +413,7 @@ def GetPoseandCostsF( inputs, outputs, ) - for ind, data in zip(inds, D): + for ind, data in zip(inds, D, strict=False): db["frame" + str(ind).zfill(strwidth)] = data del D break @@ -467,13 +447,9 @@ def GetPoseandCostsS(cfg, dlc_cfg, sess, inputs, outputs, cap, nframes, shelf_pa "minimal confidence": dlc_cfg["minconfidence"], "sigma": dlc_cfg.get("sigma", 1), "PAFgraph": dlc_cfg["partaffinityfield_graph"], - "PAFinds": dlc_cfg.get( - "paf_best", np.arange(len(dlc_cfg["partaffinityfield_graph"])) - ), + "PAFinds": dlc_cfg.get("paf_best", np.arange(len(dlc_cfg["partaffinityfield_graph"]))), "all_joints": [[i] for i in range(len(dlc_cfg["all_joints"]))], - "all_joints_names": [ - dlc_cfg["all_joints_names"][i] for i in range(len(dlc_cfg["all_joints"])) - ], + "all_joints_names": [dlc_cfg["all_joints_names"][i] for i in range(len(dlc_cfg["all_joints"]))], "nframes": nframes, } pbar = tqdm(total=nframes)