@renamed_parameter(old="videotype", new="video_extensions", since="3.0.0")
def convert_detections2tracklets(
config: str,
videos: str | list[str],
video_extensions: str | Sequence[str] | None = None,
shuffle: int = 1,
trainingsetindex: int = 0,
overwrite: bool = False,
destfolder: str | None = None,
ignore_bodyparts: list[str] | None = None,
inferencecfg: dict | None = None,
modelprefix="",
greedy: bool = False, # TODO(niels): implement greedy assembly during video analysis
calibrate: bool = False, # TODO(niels): implement assembly calibration during video analysis
window_size: int = 0, # TODO(niels): implement window size selection for assembly during video analysis
identity_only=False,
track_method="",
snapshot_index: int | str | None = None,
detector_snapshot_index: int | str | None = None,
):
"""TODO: Documentation, clean & remove code duplication (with analyze video)"""
cfg = auxiliaryfunctions.read_config(config)
inference_cfg = inferencecfg
track_method = auxfun_multianimal.get_track_method(cfg, track_method=track_method)
if len(cfg["multianimalbodyparts"]) == 1 and track_method != "box":
warnings.warn("Switching to `box` tracker for single point tracking...", stacklevel=2)
track_method = "box"
cfg["default_track_method"] = track_method
auxiliaryfunctions.write_config(config, cfg)
train_fraction = cfg["TrainingFraction"][trainingsetindex]
start_path = os.getcwd() # record cwd to return to this directory in the end
# TODO: add cropping as in video analysis!
# if cropping is not None:
# cfg['cropping']=True
# cfg['x1'],cfg['x2'],cfg['y1'],cfg['y2']=cropping
# print("Overwriting cropping parameters:", cropping)
# print("These are used for all videos, but won't be save to the cfg file.")
rel_model_dir = auxiliaryfunctions.get_model_folder(
train_fraction,
shuffle,
cfg,
modelprefix=modelprefix,
engine=Engine.PYTORCH,
)
model_dir = Path(cfg["project_path"]) / rel_model_dir
path_test_config = model_dir / "test" / "pose_cfg.yaml"
dlc_cfg = auxiliaryfunctions.read_plainconfig(str(path_test_config))
if "multi-animal" not in dlc_cfg["dataset_type"]:
raise ValueError("This function is only required for multianimal projects!")
if track_method == "ctd":
raise ValueError(
"CTD tracking occurs directly during video analysis. No need to call "
"`convert_detections2tracklets` with `track_method=='ctd'`."
)
if inference_cfg is None:
inference_cfg = auxfun_multianimal.read_inferencecfg(model_dir / "test" / "inference_cfg.yaml", cfg)
auxfun_multianimal.check_inferencecfg_sanity(cfg, inference_cfg)
if len(cfg["multianimalbodyparts"]) == 1 and track_method != "box":
warnings.warn("Switching to `box` tracker for single point tracking...", stacklevel=2)
track_method = "box"
# Also ensure `boundingboxslack` is greater than zero, otherwise overlap
# between trackers cannot be evaluated, resulting in empty tracklets.
inference_cfg["boundingboxslack"] = max(inference_cfg["boundingboxslack"], 40)
loader = DLCLoader(
config,
trainset_index=trainingsetindex,
shuffle=shuffle,
modelprefix=modelprefix,
)
snapshot_index, detector_snapshot_index = parse_snapshot_index_for_analysis(
loader.project_cfg,
loader.model_cfg,
snapshot_index,
detector_snapshot_index,
)
dlc_scorer = get_scorer_name(
cfg,
shuffle,
train_fraction,
snapshot_index=snapshot_index,
detector_index=detector_snapshot_index,
modelprefix=modelprefix,
)
paths_input = videos
videos = collect_video_paths(videos, extensions=video_extensions)
if len(videos) == 0:
print(f"No videos were found in {paths_input}")
return
for video in videos:
print("Processing... ", video)
if destfolder is None:
output_path = video.parent
else:
output_path = Path(destfolder)
output_path.mkdir(exist_ok=True, parents=True)
video_name = video.stem
data_prefix = video_name + dlc_scorer
data_filename = output_path / (data_prefix + ".h5")
print(f"Loading From {data_filename}")
data, metadata = auxfun_multianimal.LoadFullMultiAnimalData(str(data_filename))
if track_method == "ellipse":
method = "el"
elif track_method == "box":
method = "bx"
else:
method = "sk"
track_filename = output_path / (data_prefix + f"_{method}.pickle")
if not overwrite and track_filename.exists():
# TODO: check if metadata are identical (same parameters!)
print(f"Tracklets already computed at {track_filename}")
print("Set overwrite = True to overwrite.")
else:
assemblies_path = data_filename.with_stem(data_filename.stem + "_assemblies").with_suffix(".pickle")
if not assemblies_path.exists():
raise FileNotFoundError(
f"Could not find the assembles file {assemblies_path}. You're "
f"converting detections to tracklets using PyTorch, which "
"means the assemblies file must be created by the model when "
"analyzing the video!"
)
assemblies_data = auxiliaryfunctions.read_pickle(assemblies_path)
tracklets = build_tracklets(
assemblies_data=assemblies_data,
track_method=track_method,
inference_cfg=inference_cfg,
joints=data["metadata"]["all_joints_names"],
scorer=metadata["data"]["Scorer"],
num_frames=data["metadata"]["nframes"],
ignore_bodyparts=ignore_bodyparts,
unique_bodyparts=cfg["uniquebodyparts"],
identity_only=identity_only,
)
with open(track_filename, "wb") as f:
pickle.dump(tracklets, f, pickle.HIGHEST_PROTOCOL)
os.chdir(str(start_path))
print(
"The tracklets were created (i.e., under the hood "
"deeplabcut.convert_detections2tracklets was run). Now you can "
"'refine_tracklets' in the GUI, or run 'deeplabcut.stitch_tracklets'."
)