diff --git a/examples/convert_file_formats.py b/examples/convert_file_formats.py index 8eb99e94b..48c246325 100644 --- a/examples/convert_file_formats.py +++ b/examples/convert_file_formats.py @@ -18,14 +18,13 @@ from movement.io import load_poses, save_poses # %% -# Define the file path +# Load the dataset # -------------------- # This should the location to your outputfile by one of our # supported pose estimation # frameworks (e.g., DeepLabCut, SLEAP), containing predicted pose tracks. # For example, the path could be something like: -# uncomment and edit the following line to point to your own local file # file_path = "/path/to/my/data.h5" # %% @@ -36,6 +35,12 @@ "SLEAP_three-mice_Aeon_proofread.analysis.h5" )["poses"] print(fpath) + +# %% +# Now we are loading this dataset into our xarray dataset that we can then +# modify to our liking. +ds = load_poses.from_sleap_file(fpath, fps=30) +print(ds) # %% # Rename Keypoints of your choice # -------------------------------- @@ -53,41 +58,21 @@ # %% # Now we can run the function and see that -# the keypoints have been renamed -def rename_keypoints(fpath, rename_dict): - fpath = pathlib.Path(fpath) - - # this determines the filename of your modified file - # change it if you like to change the filename - dest_path = fpath.parent / f"{fpath.stem}_renamed.csv" - - if dest_path.exists(): - print(f"Skipping {fpath} as {dest_path} already exists.") - return - - if fpath.exists(): - # load the data - ds = load_poses.from_sleap_file(fpath, fps=60) - # get the current names of the keypoints - keypoint_names = ds.coords["keypoints"].values - print("Original keypoints:", keypoint_names) - # rename the keypoints - if rename_dict is None or len(rename_dict) == 0: - print("No KPs to rename. Skipping renaming step.") - else: - new_keypoints = [rename_dict.get(kp, kp) for kp in keypoint_names] - print("New keypoints:", new_keypoints) - # Assign the modified values back to the Dataset - ds = ds.assign_coords(keypoints=new_keypoints) - # save poses to dlc file format - save_poses.to_dlc_file(ds, dest_path) - # raise ValueError if path does not exist. - +# the keypoints have been renamed. +# this function takes the dataset and the rename_dict as input. +def rename_keypoints(ds, rename_dict): + # get the current names of the keypoints + keypoint_names = ds.coords["keypoints"].values + print("Original keypoints:", keypoint_names) + # rename the keypoints + if rename_dict is None or len(rename_dict) == 0: + print("No KPs to rename. Skipping renaming step.") else: - raise ValueError( - f"File '{fpath}' does not exist. " - f"Please check the file path and try again." - ) + new_keypoints = [rename_dict.get(kp, kp) for kp in keypoint_names] + print("New keypoints:", new_keypoints) + # Assign the modified values back to the Dataset + ds = ds.assign_coords(keypoints=new_keypoints) + return ds # %% @@ -95,39 +80,20 @@ def rename_keypoints(fpath, rename_dict): # -------------------------------- # First, create a list of keypoints. # to delete modify this list accordingly -kps_to_delete = ["abdomen_pre", "abdomen", "tailbase", "front_L", "front_R"] -# %% -# Now we can go ahead and delete those Keypoints - -def delete_keypoints(fpath, delete_keypoints): - fpath = pathlib.Path(fpath) +kps_to_delete = ["abdomen_pre", "abdomen", "tailbase", "front_L", "front_R"] - # this determines the filename of your modified file - # change it if you like to change the filename - dest_path = fpath.parent / f"{fpath.stem}_deletedKPs.csv" - # check if the file already exists - if dest_path.exists(): - print(f"Skipping {fpath} as {dest_path} already exists.") - return - if fpath.exists(): - # load the data - ds = load_poses.from_sleap_file(fpath, fps=60) - if delete_keypoints is None or len(delete_keypoints) == 0: - print("No KPs to delete. Skipping deleting step.") - else: - # Delete the specified keypoints - # and their corresponding data - ds = ds.drop_sel(keypoints=delete_keypoints) - # save poses to dlc file format - save_poses.to_dlc_file(ds, dest_path) - # raise ValueError if path does not exist. +# %% +# Now we can go ahead and delete those Keypoints +def delete_keypoints(ds, delete_keypoints): + if delete_keypoints is None or len(delete_keypoints) == 0: + print("No KPs to delete. Skipping deleting step.") else: - raise ValueError( - f"File '{fpath}' does not exist. " - f"Please check the file path and try again." - ) + # Delete the specified keypoints + # and their corresponding data + ds = ds.drop_sel(keypoints=delete_keypoints) + return ds # %% @@ -135,38 +101,69 @@ def delete_keypoints(fpath, delete_keypoints): # -------------------------------- # Again create a list with the # desired order of the keypoints + ordered_keypoints = ["nose", "earL", "earR", "neck", "hipL", "hipR", "tail"] + # %% # Now we can go ahead and reorder # those keypoints +def reorder_keypoints(ds, ordered_keypoints): + # reorder the keypoints + if ordered_keypoints is None or len(ordered_keypoints) == 0: + print("No KPs to reorder. Skipping reordering step.") + else: + # Reorder the keypoints in the Dataset + ds = ds.reindex(keypoints=ordered_keypoints) + return ds -def reorder_keypoints(fpath, ordered_keypoints): - fpath = pathlib.Path(fpath) - - dest_path = fpath.parent / f"{fpath.stem}_modified.csv" - - if dest_path.exists(): - print(f"Skipping {fpath} as {dest_path} already exists.") - return +# %% +# and how can I use this when I only have my filepath???? +# ----------- +# Great! Now we know how we can modify our movement datasets! +# let's put it all together and see how we could use this in a real world +# scenario. +# This function will convert all files in the folder to the desired format. +# The function will rename the keypoints, delete specified keypoints, +# and reorder them in the dataset. +# The function will then save the modified dataset to a new file. + + +FPATH = "/path/to/your/data/" + + +def convert_all(FPATH, ext="_inference.slp"): + source_folder = pathlib.Path(FPATH) + fpaths = list(source_folder.rglob(f"*{ext}")) + + for fpath in fpaths: + fpath = pathlib.Path(fpath) + + # this determines the filename of your modified file + # change it if you like to change the filename + dest_path = fpath.parent / "tracking_2D_8KP.csv" + + if dest_path.exists(): + print(f"Skipping {fpath} as {dest_path} already exists.") + return + + if fpath.exists(): + print(f"processing: {fpath}") + # load the data + ds = load_poses.from_sleap_file(fpath, fps=60) + ds_renamed = rename_keypoints(ds, rename_dict) + ds_deleted = delete_keypoints(ds_renamed, kps_to_delete) + ds_reordered = reorder_keypoints(ds_deleted, ordered_keypoints) + # save poses to dlc file format + # here we are also splitting the multi animal file into 2 + # separate files for each animal + save_poses.to_dlc_file( + ds_reordered, dest_path, split_individuals=True + ) - if fpath.exists(): - # load the data - ds = load_poses.from_sleap_file(fpath, fps=60) - # reorder the keypoints - if ordered_keypoints is None or len(ordered_keypoints) == 0: - print("No KPs to reorder. Skipping reordering step.") else: - # Reorder the keypoints in the Dataset - ds = ds.reindex(keypoints=ordered_keypoints) - - # save poses to dlc file format - save_poses.to_dlc_file(ds, dest_path) - - # raise ValueError if path does not exist. - else: - raise ValueError( - f"File '{fpath}' does not exist. " - f"Please check the file path and try again." - ) + raise ValueError( + f"File '{fpath}' does not exist. " + f"Please check the file path and try again." + )