diff --git a/deeplabcut/pose_estimation_tensorflow/predict_multianimal.py b/deeplabcut/pose_estimation_tensorflow/predict_multianimal.py index a8b7248c95..00c23b24dd 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 @@ -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 from deeplabcut.utils.auxfun_videos import VideoWriter @@ -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: 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( @@ -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) + 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"]: + 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): @@ -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 @@ -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"]) @@ -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 @@ -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 @@ -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), @@ -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( @@ -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