From 9ea18e38f9bd1b0b5c81cef6b411e0077c6d3a17 Mon Sep 17 00:00:00 2001 From: Poruri Sai Rahul Date: Fri, 13 Oct 2023 12:55:13 +0530 Subject: [PATCH 01/11] "deeplabcut.filterpredictions" and "deeplabcut.analyze_skeleton" return mappings (#1908) * REF: Denest filterpredictions for loop there were two try/excepts, one within the other. this commit simplifies the code by pulling the inner try/except out. the code now tries to look for filtered data. if filtered data is found for the video, we continue to the next item in the loop. if not found, we catch and ignore the error. we then try to load the data and filter it modified: deeplabcut/post_processing/filtering.py * REF: Reduce the code within the try block the second try block in filterpredictions looked for the analyzed file and performed operations to filter it - and if the file is not found, continue to the next item in the for loop. this commit reduces the try block to only load the analyzed file and immediately catch the error and continue if the file doesnt exist. the rest of the filtering on the analyzed data is performed outside the try/except block. modified: deeplabcut/post_processing/filtering.py * FEAT: Return video filename to filtered dataframe mapping the filterpredictions now returns a mapping instead of returning None. modified: deeplabcut/post_processing/filtering.py * REF: Reduce scope of try/except instead of finding and analyzing the skeleton in the try block, this commit reduces the scope of the try block to only finding the file and moves the skeleton calculations to outside the try/except. modified: deeplabcut/post_processing/analyze_skeleton.py * FIX: Actually return the mapping modified: deeplabcut/post_processing/filtering.py * FEAT: Return mapping from video filename to skeleton dataframe the analyze_skeleton function now returns a mapping instead of None modified: deeplabcut/post_processing/analyze_skeleton.py * Add return_data flag --------- Co-authored-by: Jessy Lauer <30733203+jeylau@users.noreply.github.com> --- .../post_processing/analyze_skeleton.py | 65 +++++--- deeplabcut/post_processing/filtering.py | 147 ++++++++++-------- 2 files changed, 127 insertions(+), 85 deletions(-) diff --git a/deeplabcut/post_processing/analyze_skeleton.py b/deeplabcut/post_processing/analyze_skeleton.py index d30587c30b..172f0c9f26 100644 --- a/deeplabcut/post_processing/analyze_skeleton.py +++ b/deeplabcut/post_processing/analyze_skeleton.py @@ -180,6 +180,7 @@ def analyzeskeleton( destfolder=None, modelprefix="", track_method="", + return_data=False, ): """Extracts length and orientation of each "bone" of the skeleton. @@ -231,15 +232,25 @@ def analyzeskeleton( For multiple animals, must be either 'box', 'skeleton', or 'ellipse' and will be taken from the config.yaml file if none is given. + return_data: bool, optional, default=False + If True, returns a dictionary of the filtered data keyed by video names. + Returns ------- - None + video_to_skeleton_df + Dictionary mapping video filepaths to skeleton dataframes. + + * If no videos exist, the dictionary will be empty. + * If a video is not analyzed, the corresponding value in the dictionary will be + None. """ # Load config file, scorer and videos cfg = auxiliaryfunctions.read_config(config) if not cfg["skeleton"]: raise ValueError("No skeleton defined in the config.yaml.") + video_to_skeleton_df = {} + track_method = auxfun_multianimal.get_track_method(cfg, track_method=track_method) DLCscorer, DLCscorerlegacy = auxiliaryfunctions.get_scorer_name( cfg, @@ -259,33 +270,39 @@ def analyzeskeleton( df, filepath, scorer, _ = auxiliaryfunctions.load_analyzed_data( destfolder, vname, DLCscorer, filtered, track_method ) - output_name = filepath.replace(".h5", f"_skeleton.h5") - if os.path.isfile(output_name): - print(f"Skeleton in video {vname} already processed. Skipping...") - continue - - bones = {} - if "individuals" in df.columns.names: - for animal_name, df_ in df.groupby(level="individuals", axis=1): - temp = df_.droplevel(["scorer", "individuals"], axis=1) - if animal_name != "single": - for bp1, bp2 in cfg["skeleton"]: - name = "{}_{}_{}".format(animal_name, bp1, bp2) - bones[name] = analyzebone(temp[bp1], temp[bp2]) - else: - for bp1, bp2 in cfg["skeleton"]: - name = "{}_{}".format(bp1, bp2) - bones[name] = analyzebone(df[scorer][bp1], df[scorer][bp2]) - - skeleton = pd.concat(bones, axis=1) - skeleton.to_hdf(output_name, "df_with_missing", format="table", mode="w") - if save_as_csv: - skeleton.to_csv(output_name.replace(".h5", ".csv")) - except FileNotFoundError as e: print(e) + video_to_skeleton_df[video] = None continue + output_name = filepath.replace(".h5", f"_skeleton.h5") + if os.path.isfile(output_name): + print(f"Skeleton in video {vname} already processed. Skipping...") + video_to_skeleton_df[video] = pd.read_hdf(output_name, "df_with_missing") + continue + + bones = {} + if "individuals" in df.columns.names: + for animal_name, df_ in df.groupby(level="individuals", axis=1): + temp = df_.droplevel(["scorer", "individuals"], axis=1) + if animal_name != "single": + for bp1, bp2 in cfg["skeleton"]: + name = "{}_{}_{}".format(animal_name, bp1, bp2) + bones[name] = analyzebone(temp[bp1], temp[bp2]) + else: + for bp1, bp2 in cfg["skeleton"]: + name = "{}_{}".format(bp1, bp2) + bones[name] = analyzebone(df[scorer][bp1], df[scorer][bp2]) + + skeleton = pd.concat(bones, axis=1) + video_to_skeleton_df[video] = skeleton + skeleton.to_hdf(output_name, "df_with_missing", format="table", mode="w") + if save_as_csv: + skeleton.to_csv(output_name.replace(".h5", ".csv")) + + if return_data: + return video_to_skeleton_df + if __name__ == "__main__": parser = argparse.ArgumentParser() diff --git a/deeplabcut/post_processing/filtering.py b/deeplabcut/post_processing/filtering.py index 1c81081b6b..d462bbd249 100644 --- a/deeplabcut/post_processing/filtering.py +++ b/deeplabcut/post_processing/filtering.py @@ -82,6 +82,7 @@ def filterpredictions( destfolder=None, modelprefix="", track_method="", + return_data=False, ): """Fits frame-by-frame pose predictions. @@ -148,9 +149,17 @@ def filterpredictions( For multiple animals, must be either 'box', 'skeleton', or 'ellipse' and will be taken from the config.yaml file if none is given. + return_data: bool, optional, default=False + If True, returns a dictionary of the filtered data keyed by video names. + Returns ------- - None + video_to_filtered_df + Dictionary mapping video filepaths to filtered dataframes. + + * If no videos exist, the dictionary will be empty. + * If a video is not analyzed, the corresponding value in the dictionary will be + None. Examples -------- @@ -202,9 +211,12 @@ def filterpredictions( ) Videos = auxiliaryfunctions.get_list_of_videos(video, videotype) + video_to_filtered_df = {} + if not len(Videos): print("No video(s) were found. Please check your paths and/or 'videotype'.") - return + if return_data: + return video_to_filtered_df for video in Videos: if destfolder is None: @@ -214,69 +226,82 @@ def filterpredictions( vname = Path(video).stem try: - _ = auxiliaryfunctions.load_analyzed_data( + df, filepath, _, _ = auxiliaryfunctions.load_analyzed_data( destfolder, vname, DLCscorer, True, track_method ) print(f"Data from {vname} were already filtered. Skipping...") - except FileNotFoundError: # Data haven't been filtered yet - try: - df, filepath, _, _ = auxiliaryfunctions.load_analyzed_data( - destfolder, vname, DLCscorer, track_method=track_method + video_to_filtered_df[video] = df + # Data has been filtered so continue to the next video + continue + except FileNotFoundError: + pass + + # Data haven't been filtered yet + try: + df, filepath, _, _ = auxiliaryfunctions.load_analyzed_data( + destfolder, vname, DLCscorer, track_method=track_method + ) + except FileNotFoundError as e: + video_to_filtered_df[video] = None + print(e) + continue + + nrows = df.shape[0] + if filtertype == "arima": + temp = df.values.reshape((nrows, -1, 3)) + placeholder = np.empty_like(temp) + for i in range(temp.shape[1]): + x, y, p = temp[:, i].T + meanx, _ = FitSARIMAXModel( + x, p, p_bound, alpha, ARdegree, MAdegree, False + ) + meany, _ = FitSARIMAXModel( + y, p, p_bound, alpha, ARdegree, MAdegree, False ) - nrows = df.shape[0] - if filtertype == "arima": - temp = df.values.reshape((nrows, -1, 3)) - placeholder = np.empty_like(temp) - for i in range(temp.shape[1]): - x, y, p = temp[:, i].T - meanx, _ = FitSARIMAXModel( - x, p, p_bound, alpha, ARdegree, MAdegree, False - ) - meany, _ = FitSARIMAXModel( - y, p, p_bound, alpha, ARdegree, MAdegree, False - ) - meanx[0] = x[0] - meany[0] = y[0] - placeholder[:, i] = np.c_[meanx, meany, p] - data = pd.DataFrame( - placeholder.reshape((nrows, -1)), - columns=df.columns, - index=df.index, - ) - elif filtertype == "median": - data = df.copy() - mask = data.columns.get_level_values("coords") != "likelihood" - data.loc[:, mask] = df.loc[:, mask].apply( - signal.medfilt, args=(windowlength,), axis=0 - ) - elif filtertype == "spline": - data = df.copy() - mask_data = data.columns.get_level_values("coords").isin(("x", "y")) - xy = data.loc[:, mask_data].values - prob = data.loc[:, ~mask_data].values - missing = np.isnan(xy) - xy_filled = columnwise_spline_interp(xy, windowlength) - filled = ~np.isnan(xy_filled) - xy[filled] = xy_filled[filled] - inds = np.argwhere(missing & filled) - if inds.size: - # Retrieve original individual label indices - inds[:, 1] //= 2 - inds = np.unique(inds, axis=0) - prob[inds[:, 0], inds[:, 1]] = 0.01 - data.loc[:, ~mask_data] = prob - data.loc[:, mask_data] = xy - else: - raise ValueError(f"Unknown filter type {filtertype}") - - outdataname = filepath.replace(".h5", "_filtered.h5") - data.to_hdf(outdataname, "df_with_missing", format="table", mode="w") - if save_as_csv: - print("Saving filtered csv poses!") - data.to_csv(outdataname.split(".h5")[0] + ".csv") - except FileNotFoundError as e: - print(e) - continue + meanx[0] = x[0] + meany[0] = y[0] + placeholder[:, i] = np.c_[meanx, meany, p] + data = pd.DataFrame( + placeholder.reshape((nrows, -1)), + columns=df.columns, + index=df.index, + ) + elif filtertype == "median": + data = df.copy() + mask = data.columns.get_level_values("coords") != "likelihood" + data.loc[:, mask] = df.loc[:, mask].apply( + signal.medfilt, args=(windowlength,), axis=0 + ) + elif filtertype == "spline": + data = df.copy() + mask_data = data.columns.get_level_values("coords").isin(("x", "y")) + xy = data.loc[:, mask_data].values + prob = data.loc[:, ~mask_data].values + missing = np.isnan(xy) + xy_filled = columnwise_spline_interp(xy, windowlength) + filled = ~np.isnan(xy_filled) + xy[filled] = xy_filled[filled] + inds = np.argwhere(missing & filled) + if inds.size: + # Retrieve original individual label indices + inds[:, 1] //= 2 + inds = np.unique(inds, axis=0) + prob[inds[:, 0], inds[:, 1]] = 0.01 + data.loc[:, ~mask_data] = prob + data.loc[:, mask_data] = xy + else: + raise ValueError(f"Unknown filter type {filtertype}") + + video_to_filtered_df[video] = data + + outdataname = filepath.replace(".h5", "_filtered.h5") + data.to_hdf(outdataname, "df_with_missing", format="table", mode="w") + if save_as_csv: + print("Saving filtered csv poses!") + data.to_csv(outdataname.split(".h5")[0] + ".csv") + + if return_data: + return video_to_filtered_df if __name__ == "__main__": From 3615f75d95f0b3c2450135c8b82a35bde517f004 Mon Sep 17 00:00:00 2001 From: Jessy Lauer <30733203+jeylau@users.noreply.github.com> Date: Mon, 16 Oct 2023 22:14:24 +0200 Subject: [PATCH 02/11] Make search window size a free parameter (#2414) --- deeplabcut/pose_estimation_3d/camera_calibration.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/deeplabcut/pose_estimation_3d/camera_calibration.py b/deeplabcut/pose_estimation_3d/camera_calibration.py index 6733b618be..1fbb9d8b18 100644 --- a/deeplabcut/pose_estimation_3d/camera_calibration.py +++ b/deeplabcut/pose_estimation_3d/camera_calibration.py @@ -26,7 +26,7 @@ matplotlib_axes_logger.setLevel("ERROR") -def calibrate_cameras(config, cbrow=8, cbcol=6, calibrate=False, alpha=0.4): +def calibrate_cameras(config, cbrow=8, cbcol=6, calibrate=False, alpha=0.4, search_window_size=(11, 11)): """This function extracts the corners points from the calibration images, calibrates the camera and stores the calibration files in the project folder (defined in the config file). Make sure you have around 20-60 pairs of calibration images. The function should be used iteratively to select the right set of calibration images. @@ -57,6 +57,9 @@ def calibrate_cameras(config, cbrow=8, cbcol=6, calibrate=False, alpha=0.4): i.e. the rectified images are zoomed in. When alpha = 1, all the pixels from the original images are retained. For more details: https://docs.opencv.org/2.4/modules/calib3d/doc/camera_calibration_and_3d_reconstruction.html + search_window_size: tuple of int + Half of the side length of the search window when refining detected checkerboard corners for subpixel accuracy. + Example -------- Linux/MacOs/Windows @@ -139,7 +142,7 @@ def calibrate_cameras(config, cbrow=8, cbcol=6, calibrate=False, alpha=0.4): img_shape[cam] = gray.shape[::-1] objpoints[cam].append(objp) corners = cv2.cornerSubPix( - gray, corners, (11, 11), (-1, -1), criteria + gray, corners, search_window_size, (-1, -1), criteria ) imgpoints[cam].append(corners) # Draw the corners and store the images From 284b2ceaceda282f36d922ff2c39118122f0fd20 Mon Sep 17 00:00:00 2001 From: Jessy Lauer <30733203+jeylau@users.noreply.github.com> Date: Tue, 17 Oct 2023 10:41:51 +0200 Subject: [PATCH 03/11] Install DLC from pypi rather than its local version (#2416) --- conda-environments/DEEPLABCUT_M1.yaml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/conda-environments/DEEPLABCUT_M1.yaml b/conda-environments/DEEPLABCUT_M1.yaml index e66016f770..e5a8bab83d 100644 --- a/conda-environments/DEEPLABCUT_M1.yaml +++ b/conda-environments/DEEPLABCUT_M1.yaml @@ -40,4 +40,4 @@ dependencies: - ffmpeg - apple::tensorflow-deps - pip: - - -e ../[gui,apple_mchips] + - "deeplabcut[gui,apple_mchips]" From 7fa496c1e8cc9d0e4ba48ff1d5ec851eb790256b Mon Sep 17 00:00:00 2001 From: Christopher Bottoms Date: Tue, 24 Oct 2023 08:17:22 -0500 Subject: [PATCH 04/11] Correct link to platform specifications page (#2423) --- docs/UseOverviewGuide.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/UseOverviewGuide.md b/docs/UseOverviewGuide.md index 6f2d69e5df..d19c1feec5 100644 --- a/docs/UseOverviewGuide.md +++ b/docs/UseOverviewGuide.md @@ -24,7 +24,7 @@ Getting Started: [a video tutorial on navigating the documentation!](https://www - **a set of videos that span the types of behaviors you want to track.** Having 10 videos that include different backgrounds, different individuals, and different postures is MUCH better than 1 or 2 videos of 1 or 2 different individuals (i.e. 10-20 frames from each of 10 videos is **much better** than 50-100 frames from 2 videos). - - **minimally, a computer w/a CPU.** If you want to use DeepLabCut on your own computer for many experiments, then you should get an NVIDIA GPU. See technical specs [here](https://github.com/DeepLabCutDeepLabCut/wiki/FAQ). You can also use cloud computing resources, including COLAB ([see how](https://github.com/DeepLabCut/DeepLabCut/blob/master/examples/README.md)). + - **minimally, a computer w/a CPU.** If you want to use DeepLabCut on your own computer for many experiments, then you should get an NVIDIA GPU. See technical specs [here](https://github.com/DeepLabCut/DeepLabCut/wiki/FAQ). You can also use cloud computing resources, including COLAB ([see how](https://github.com/DeepLabCut/DeepLabCut/blob/master/examples/README.md)). ### What you DON'T need to get started: From f0cf278d953a1431d88ba80c7178337686a6e3ad Mon Sep 17 00:00:00 2001 From: Jessy Lauer <30733203+jeylau@users.noreply.github.com> Date: Tue, 24 Oct 2023 15:17:38 +0200 Subject: [PATCH 05/11] Automatically check for napari-deeplabcut updates (in addition to DLC) (#2422) * Also check for napari-deeplabcut updates * Increase min napari-deeplabcut required version --- deeplabcut/gui/window.py | 28 ++++++++++++++++++++-------- setup.py | 2 +- 2 files changed, 21 insertions(+), 9 deletions(-) diff --git a/deeplabcut/gui/window.py b/deeplabcut/gui/window.py index e8917115b1..71f291c73e 100644 --- a/deeplabcut/gui/window.py +++ b/deeplabcut/gui/window.py @@ -22,6 +22,7 @@ from deeplabcut.gui import BASE_DIR, components, utils from deeplabcut.gui.tabs import * from deeplabcut.gui.widgets import StreamReceiver, StreamWriter +from napari_deeplabcut import misc from PySide6.QtWidgets import QMessageBox, QMenu, QWidget, QMainWindow from PySide6 import QtCore from PySide6.QtGui import QIcon, QAction @@ -31,9 +32,25 @@ def _check_for_updates(): is_latest, latest_version = utils.is_latest_deeplabcut_version() - if not is_latest: + is_latest_plugin, latest_plugin_version = misc.is_latest_version() + if is_latest and is_latest_plugin: msg = QtWidgets.QMessageBox( - text=f"DeepLabCut {latest_version} available", + text=f"DeepLabCut is up-to-date", + ) + msg.exec_() + else: + if not is_latest and is_latest_plugin: + text = f"DeepLabCut {latest_version} available" + command = "pip", "install", "-U", "deeplabcut" + elif not is_latest_plugin and is_latest: + text = f"DeepLabCut labeling plugin {latest_plugin_version} available" + command = "pip", "install", "-U", "napari-deeplabcut" + else: + text = f"DeepLabCut {latest_version}\nand labeling plugin {latest_plugin_version} available" + command = "pip", "install", "-U", "deeplabcut", "napari-deeplabcut" + + msg = QtWidgets.QMessageBox( + text=text, ) msg.setIcon(QtWidgets.QMessageBox.Information) update_btn = msg.addButton("Update", msg.AcceptRole) @@ -42,13 +59,8 @@ def _check_for_updates(): msg.exec_() if msg.clickedButton() is update_btn: subprocess.check_call( - [sys.executable, "-m", "pip", "install", "-U", "deeplabcut"] + [sys.executable, "-m", *command] ) - else: - msg = QtWidgets.QMessageBox( - text=f"DeepLabCut is up-to-date", - ) - msg.exec_() class MainWindow(QMainWindow): diff --git a/setup.py b/setup.py index 7cb3f20253..fe2fd995bc 100644 --- a/setup.py +++ b/setup.py @@ -51,7 +51,7 @@ "gui": [ "pyside6<6.3.2", "qdarkstyle==3.1", - "napari-deeplabcut>=0.2", + "napari-deeplabcut>=0.2.1.2", ], "openvino": ["openvino-dev==2022.1.0"], "docs": ["numpydoc"], From 4de123626633c88b54be33cae665d9fe9ad9e6a2 Mon Sep 17 00:00:00 2001 From: Jessy Lauer <30733203+jeylau@users.noreply.github.com> Date: Tue, 24 Oct 2023 17:48:42 +0200 Subject: [PATCH 06/11] Smartly restore pretrained model weights (#2426) --- .../pose_estimation_tensorflow/core/train.py | 4 +++- .../core/train_multianimal.py | 4 +++- deeplabcut/utils/auxfun_models.py | 17 +++++++++++++++++ 3 files changed, 23 insertions(+), 2 deletions(-) diff --git a/deeplabcut/pose_estimation_tensorflow/core/train.py b/deeplabcut/pose_estimation_tensorflow/core/train.py index b2d72157f4..487df6b462 100644 --- a/deeplabcut/pose_estimation_tensorflow/core/train.py +++ b/deeplabcut/pose_estimation_tensorflow/core/train.py @@ -31,6 +31,7 @@ ) from deeplabcut.pose_estimation_tensorflow.nnets import PoseNetFactory from deeplabcut.pose_estimation_tensorflow.util.logging import setup_logging +from deeplabcut.utils import auxfun_models class LearningRate(object): @@ -243,7 +244,8 @@ def train( sess.run(tf.compat.v1.local_variables_initializer()) # Restore variables from disk. - restorer.restore(sess, cfg["init_weights"]) + auxfun_models.smart_restore(restorer, sess, cfg["init_weights"], net_type) + if maxiters is None: max_iter = int(cfg["multi_step"][-1][1]) else: diff --git a/deeplabcut/pose_estimation_tensorflow/core/train_multianimal.py b/deeplabcut/pose_estimation_tensorflow/core/train_multianimal.py index f9f2caed29..08ea4ce7fc 100644 --- a/deeplabcut/pose_estimation_tensorflow/core/train_multianimal.py +++ b/deeplabcut/pose_estimation_tensorflow/core/train_multianimal.py @@ -28,6 +28,7 @@ get_optimizer, LearningRate, ) +from deeplabcut.utils import auxfun_models def train( @@ -157,7 +158,8 @@ def train( sess.run(tf.compat.v1.global_variables_initializer()) sess.run(tf.compat.v1.local_variables_initializer()) - restorer.restore(sess, cfg["init_weights"]) + auxfun_models.smart_restore(restorer, sess, cfg["init_weights"], net_type) + if maxiters is None: max_iter = int(cfg["multi_step"][-1][1]) else: diff --git a/deeplabcut/utils/auxfun_models.py b/deeplabcut/utils/auxfun_models.py index 046528219f..db7bcbfe63 100644 --- a/deeplabcut/utils/auxfun_models.py +++ b/deeplabcut/utils/auxfun_models.py @@ -168,6 +168,23 @@ def set_visible_devices(gputouse: int): tf.config.set_visible_devices(physical_devices[gputouse], "GPU") +def smart_restore(restorer, sess, checkpoint_path, net_type): + "Restore pretrained weights, smartly redownloading them if missing." + try: + restorer.restore(sess, checkpoint_path) + except ValueError as e: # The path may be wrong, or the weights no longer exist + dlcparent_path = auxiliaryfunctions.get_deeplabcut_path() + correct_model_path = os.path.join( + dlcparent_path, MODELTYPE_FILEPATH_MAP[net_type], + ) + if checkpoint_path == correct_model_path: + # The path is right, hence the weights are missing; we'll download them again. + _ = check_for_weights(net_type, Path(dlcparent_path)) + restorer.restore(sess, checkpoint_path) + else: + raise ValueError(e) + + # Aliases for backwards-compatibility Check4Weights = check_for_weights Downloadweights = download_weights From 25f8c5026a03168358360e70c98039a7d7ed1b32 Mon Sep 17 00:00:00 2001 From: biol-jsh Date: Fri, 27 Oct 2023 05:14:02 -0600 Subject: [PATCH 07/11] Fix fliplr augmentation for multi animal 2 (#2043) * Fix fliplr augmentation for multi animal 2 The current implementation of fliplr breaks when keypoints are hidden since only labels for visible points are forwarded to the pipeline. This is the second attempt at fixing this as an earlier attempt failed due to the error fixed in PR #2037. * added missing batch_joints append from DLC PR 1946 * made same changes to get_batch_from_video as get_batch. style fixes. * Fix indices when training with identity * Fix and expand unit tests --------- Co-authored-by: Niels Poulsen Co-authored-by: Jessy Lauer <30733203+jeylau@users.noreply.github.com> --- .../datasets/pose_multianimal_imgaug.py | 65 +++++++++---------- tests/test_dataset_augmentation.py | 24 +++++++ tests/test_pose_multianimal_imgaug.py | 16 +++-- 3 files changed, 65 insertions(+), 40 deletions(-) diff --git a/deeplabcut/pose_estimation_tensorflow/datasets/pose_multianimal_imgaug.py b/deeplabcut/pose_estimation_tensorflow/datasets/pose_multianimal_imgaug.py index d25c2e7ffb..8659a05314 100644 --- a/deeplabcut/pose_estimation_tensorflow/datasets/pose_multianimal_imgaug.py +++ b/deeplabcut/pose_estimation_tensorflow/datasets/pose_multianimal_imgaug.py @@ -162,7 +162,7 @@ def _load_pseudo_data_from_h5(self, cfg, threshold=0.5, mask_kpts_below_thresh=F item.joints = {} if not mask_kpts_below_thresh: - joints = np.concatenate([joint_ids, kpts], axis=1) + joints = np.concatenate([joint_ids, kpts], axis=1) joints = np.nan_to_num(joints, nan=0) else: for kpt_id, kpt in enumerate(kpts): @@ -172,7 +172,7 @@ def _load_pseudo_data_from_h5(self, cfg, threshold=0.5, mask_kpts_below_thresh=F kpts[kpt_id][:-1] = -1 kpts[kpt_id][-1] = 1 joints = np.concatenate([joint_ids, kpts], axis=1) - + sparse_joints = [] for coord in joints: @@ -349,21 +349,19 @@ def get_aug_param(cfg_value): def get_batch_from_video(self): num_images = len(self.vid) - size = self.batch_size batch_images = [] batch_joints = [] joint_ids = [] - inds_visible = [] data_items = [] trim_ends = self.cfg.get('trim_ends', None) if trim_ends is None: trim_ends = 0 # because of the existence of threshold, sampling population is adjusted to len(self.data) - img_idx = np.random.choice(len(self.data) - trim_ends *2, size=self.batch_size, replace=True) + img_idx = np.random.choice(len(self.data) - trim_ends *2, size=self.batch_size, replace=True) for i in range(self.batch_size): index = img_idx[i] offset = trim_ends - data_item = self.data[index + offset] + data_item = self.data[index + offset] data_items.append(data_item) im_file = data_item.im_path @@ -371,31 +369,29 @@ def get_batch_from_video(self): self.vid.set_to_frame(index + offset) image = self.vid.read_frame() if self.has_gt: - Joints = data_item.joints - if len(Joints[0]) == 0: + joints = data_item.joints + if len(joints[0]) == 0: # empty prediction for this frame - return None, None, None, None, None - kpts = np.zeros((self._n_kpts * self._n_animals, 2)) + return None, None, None, None + + kpts = np.full((self._n_kpts * self._n_animals, 2), np.nan) for j in range(self._n_animals): - for n, x, y in Joints.get(j, []): + for n, x, y in joints.get(j, []): kpts[j * self._n_kpts + int(n)] = x, y - joint_id = [ - Joints[person_id][:, 0].astype(int) for person_id in Joints.keys() - ] + + joint_id = np.array(list(range(self._n_kpts)) * self._n_animals) joint_ids.append(joint_id) batch_joints.append(kpts) - inds_visible.append(np.flatnonzero(np.all(kpts != 0, axis=1))) batch_images.append(image) - return batch_images, joint_ids, batch_joints, inds_visible, data_items + return batch_images, joint_ids, batch_joints, data_items def get_batch(self): img_idx = np.random.choice(self.num_images, size=self.batch_size, replace=True) batch_images = [] batch_joints = [] joint_ids = [] - inds_visible = [] data_items = [] for i in range(self.batch_size): data_item = self.data[img_idx[i]] @@ -408,21 +404,21 @@ def get_batch(self): os.path.join(self.cfg["project_path"], im_file), mode="skimage" ) if self.has_gt: - Joints = data_item.joints - kpts = np.zeros((self._n_kpts * self._n_animals, 2)) + joints = data_item.joints + kpts = np.full((self._n_kpts * self._n_animals, 2), np.nan) for j in range(self._n_animals): - for n, x, y in Joints.get(j, []): + for n, x, y in joints.get(j, []): kpts[j * self._n_kpts + int(n)] = x, y joint_id = [ - Joints[person_id][:, 0].astype(int) for person_id in Joints.keys() + np.array(list(range(self._n_kpts))) + for _ in range(self._n_animals) ] joint_ids.append(joint_id) batch_joints.append(kpts) - inds_visible.append(np.flatnonzero(np.all(kpts != 0, axis=1))) batch_images.append(image) - return batch_images, joint_ids, batch_joints, inds_visible, data_items + return batch_images, joint_ids, batch_joints, data_items def get_targetmaps_update( self, @@ -498,17 +494,11 @@ def next_batch(self, plotting=False): batch_images, joint_ids, batch_joints, - inds_visible, data_items, ) = self.get_batch_from_video() else: - ( - batch_images, - joint_ids, - batch_joints, - inds_visible, - data_items, - ) = self.get_batch() + batch_images, joint_ids, batch_joints, data_items = self.get_batch() + # in case it's empty prediction if batch_joints is None or batch_images is None: continue @@ -525,8 +515,10 @@ def next_batch(self, plotting=False): # Discard keypoints whose coordinates lie outside the cropped image batch_joints_valid = [] joint_ids_valid = [] - for joints, ids, visible in zip(batch_joints, joint_ids, inds_visible): - joints = joints[visible] + + for joints, ids in zip(batch_joints, joint_ids): + # Invisible joints are represented by nans + visible = ~np.isnan(joints[:, 0]) inside = np.logical_and.reduce( ( joints[:, 0] < image_shape[1], @@ -535,12 +527,15 @@ def next_batch(self, plotting=False): joints[:, 1] > 0, ) ) - batch_joints_valid.append(joints[inside]) + mask = visible & inside + batch_joints_valid.append(joints[mask]) + temp = [] start = 0 for array in ids: end = start + array.size - temp.append(array[inside[start:end]]) + inds = np.arange(start, end) + temp.append(array[mask[inds]]) start = end joint_ids_valid.append(temp) diff --git a/tests/test_dataset_augmentation.py b/tests/test_dataset_augmentation.py index 963bf2f541..9dab238b1e 100644 --- a/tests/test_dataset_augmentation.py +++ b/tests/test_dataset_augmentation.py @@ -109,3 +109,27 @@ def test_keypoint_horizontal_flip( temp[:, pair] = temp[:, pair[::-1]] keypoints_unaug = temp.reshape((-1, 2)) np.testing.assert_allclose(keypoints_unaug, keypoints_flipped) + + +def test_keypoint_horizontal_flip_with_nans( + sample_image, + sample_keypoints, +): + sample_keypoints[::12] = np.nan + sample_keypoints[2::12] = np.nan + keypoints_flipped = sample_keypoints.copy() + keypoints_flipped[:, 0] = sample_image.shape[1] - keypoints_flipped[:, 0] + pairs = [(0, 1), (2, 3)] + aug = augmentation.KeypointFliplr( + keypoints=list(map(str, range(12))), + symmetric_pairs=pairs, + ) + keypoints_aug = aug( + images=[sample_image], + keypoints=[sample_keypoints], + )[1][0] + temp = keypoints_aug.reshape((3, 12, 2)) + for pair in pairs: + temp[:, pair] = temp[:, pair[::-1]] + keypoints_unaug = temp.reshape((-1, 2)) + np.testing.assert_allclose(keypoints_unaug, keypoints_flipped) diff --git a/tests/test_pose_multianimal_imgaug.py b/tests/test_pose_multianimal_imgaug.py index 0f19133875..c0c8a8c5b0 100644 --- a/tests/test_pose_multianimal_imgaug.py +++ b/tests/test_pose_multianimal_imgaug.py @@ -66,7 +66,7 @@ def test_calc_target_and_scoremap_sizes( def test_get_batch(ma_dataset): for batch_size in 1, 4, 8, 16: ma_dataset.batch_size = batch_size - batch_images, joint_ids, batch_joints, _, data_items = ma_dataset.get_batch() + batch_images, joint_ids, batch_joints, data_items = ma_dataset.get_batch() assert ( len(batch_images) == len(joint_ids) @@ -74,10 +74,17 @@ def test_get_batch(ma_dataset): == len(data_items) == batch_size ) - for data_item, joint_id in zip(data_items, joint_ids): + for data_item, joint_id, batch_joint in zip(data_items, joint_ids, batch_joints): assert len(data_item.joints) == len(joint_id) + assert len(batch_joint) == len(np.concatenate(joint_id)) + start = 0 + mask = ~np.isnan(batch_joint).any(axis=1) for joints, id_ in zip(data_item.joints.values(), joint_id): - np.testing.assert_equal(joints[:, 0], id_) + inds = id_ + start + mask_ = mask[inds] + np.testing.assert_equal(joints[:, 0], id_[mask_]) + np.testing.assert_equal(joints[:, 1:], batch_joint[inds][mask_]) + start += id_.size def test_build_augmentation_pipeline(ma_dataset): @@ -88,8 +95,7 @@ def test_build_augmentation_pipeline(ma_dataset): @pytest.mark.parametrize("num_idchannel", range(4)) def test_get_targetmaps(ma_dataset, num_idchannel): ma_dataset.cfg["num_idchannel"] = num_idchannel - batch = list(ma_dataset.get_batch()[1:]) - batch.pop(2) + batch = ma_dataset.get_batch()[1:] target_size, sm_size = ma_dataset.calc_target_and_scoremap_sizes() scale = np.mean(target_size / ma_dataset.default_size) maps = ma_dataset.get_targetmaps_update(*batch, sm_size, scale) From 11f12d84406cf4e249fc618365336c364537dfa8 Mon Sep 17 00:00:00 2001 From: shaokai Date: Mon, 6 Nov 2023 15:15:43 +0100 Subject: [PATCH 08/11] Shaokai/sa transfer learning (#2389) * Added SA finetune and corresponding project --------- Co-authored-by: Jessy Lauer <30733203+jeylau@users.noreply.github.com> Co-authored-by: Timokleia <86098649+Timokleia@users.noreply.github.com> Co-authored-by: Mackenzie Mathis --- .../frame_extraction.py | 1 - .../trainingsetmanipulation.py | 35 ++- deeplabcut/gui/tabs/extract_frames.py | 6 +- deeplabcut/gui/tracklet_toolbox.py | 4 +- deeplabcut/gui/window.py | 8 +- .../modelzoo/api/spatiotemporal_adapt.py | 5 +- .../modelzoo/api/superanimal_inference.py | 2 +- .../core/train_multianimal.py | 12 +- .../datasets/pose_deterministic.py | 4 +- .../datasets/pose_imgaug.py | 8 +- .../datasets/pose_multianimal_imgaug.py | 20 +- .../datasets/pose_tensorpack.py | 4 +- .../lib/inferenceutils.py | 6 +- .../pose_estimation_tensorflow/nnets/utils.py | 4 +- .../predict_supermodel.py | 23 +- .../predict_videos.py | 12 +- .../pose_estimation_tensorflow/training.py | 219 +++++++++++------- .../util/visualize.py | 2 +- .../visualizemaps.py | 2 +- .../model/backbones/vit_pytorch.py | 2 +- .../pose_tracking_pytorch/solver/cosine_lr.py | 8 +- .../tracking_utils/reranking.py | 4 +- .../refine_training_dataset/outlier_frames.py | 6 +- deeplabcut/refine_training_dataset/stitch.py | 12 +- deeplabcut/utils/make_labeled_video.py | 6 +- docs/ModelZoo.md | 28 ++- ...estscript_superanimal_transfer_learning.py | 31 +++ setup.py | 6 +- tests/test_dataset_augmentation.py | 5 +- tests/test_inferenceutils.py | 2 +- 30 files changed, 331 insertions(+), 156 deletions(-) create mode 100644 examples/testscript_superanimal_transfer_learning.py diff --git a/deeplabcut/generate_training_dataset/frame_extraction.py b/deeplabcut/generate_training_dataset/frame_extraction.py index 22b029263e..6264e01157 100755 --- a/deeplabcut/generate_training_dataset/frame_extraction.py +++ b/deeplabcut/generate_training_dataset/frame_extraction.py @@ -262,7 +262,6 @@ def extract_frames( from deeplabcut.utils import frameselectiontools from deeplabcut.utils import auxiliaryfunctions - config_file = Path(config).resolve() cfg = auxiliaryfunctions.read_config(config_file) print("Config file read successfully.") diff --git a/deeplabcut/generate_training_dataset/trainingsetmanipulation.py b/deeplabcut/generate_training_dataset/trainingsetmanipulation.py index 2b292a54ba..6d237c27b3 100755 --- a/deeplabcut/generate_training_dataset/trainingsetmanipulation.py +++ b/deeplabcut/generate_training_dataset/trainingsetmanipulation.py @@ -31,6 +31,8 @@ auxfun_multianimal, ) from deeplabcut.utils.auxfun_videos import VideoReader +from deeplabcut.pose_estimation_tensorflow.config import load_config +from deeplabcut.modelzoo.utils import parse_available_supermodels def comparevideolistsanddatafolders(config): @@ -482,7 +484,9 @@ def merge_annotateddatasets(cfg, trainingsetfolder_full): data = pd.read_hdf(file_path) conversioncode.guarantee_multiindex_rows(data) if data.columns.levels[0][0] != cfg["scorer"]: - print(f"{file_path} labeled by a different scorer. This data will not be utilized in training dataset creation. If you need to merge datasets across scorers, see https://github.com/DeepLabCut/DeepLabCut/wiki/Using-labeled-data-in-DeepLabCut-that-was-annotated-elsewhere-(or-merge-across-labelers)") + print( + f"{file_path} labeled by a different scorer. This data will not be utilized in training dataset creation. If you need to merge datasets across scorers, see https://github.com/DeepLabCut/DeepLabCut/wiki/Using-labeled-data-in-DeepLabCut-that-was-annotated-elsewhere-(or-merge-across-labelers)" + ) continue AnnotationData.append(data) except FileNotFoundError: @@ -728,6 +732,7 @@ def create_training_dataset( net_type=None, augmenter_type=None, posecfg_template=None, + superanimal_name="", ): """Creates a training dataset. @@ -791,6 +796,10 @@ def create_training_dataset( parameters a previous training iteration. None uses the default ``pose_cfg.yaml``. + superanimal_name: string, optional, default="" + Specify the superanimal name is transfer learning with superanimal is desired. This makes sure the pose config template uses superanimal configs as template + + Returns ------- list(tuple) or None @@ -832,8 +841,23 @@ def create_training_dataset( # Loading metadata from config file: cfg = auxiliaryfunctions.read_config(config) + dlc_root_path = auxiliaryfunctions.get_deeplabcut_path() + + if superanimal_name != "": + supermodels = parse_available_supermodels() + posecfg_template = os.path.join( + dlc_root_path, + "pose_estimation_tensorflow", + "superanimal_configs", + supermodels[superanimal_name], + ) + if posecfg_template: - if not posecfg_template.endswith("pose_cfg.yaml"): + if ( + not posecfg_template.endswith("pose_cfg.yaml") + and not posecfg_template.endswith("superquadruped.yaml") + and not posecfg_template.endswith("supertopview.yaml") + ): raise ValueError( "posecfg_template argument must contain path to a pose_cfg.yaml file" ) @@ -841,14 +865,16 @@ def create_training_dataset( print("Reloading pose_cfg parameters from " + posecfg_template + "\n") from deeplabcut.utils.auxiliaryfunctions import read_plainconfig - prior_cfg = read_plainconfig(posecfg_template) + prior_cfg = read_plainconfig(posecfg_template) if cfg.get("multianimalproject", False): from deeplabcut.generate_training_dataset.multiple_individuals_trainingsetmanipulation import ( create_multianimaltraining_dataset, ) create_multianimaltraining_dataset( - config, num_shuffles, Shuffles, + config, + num_shuffles, + Shuffles, net_type=net_type, trainIndices=trainIndices, testIndices=testIndices, @@ -880,6 +906,7 @@ def create_training_dataset( "resnet" in net_type or "mobilenet" in net_type or "efficientnet" in net_type + or "dlcrnet" in net_type ): pass else: diff --git a/deeplabcut/gui/tabs/extract_frames.py b/deeplabcut/gui/tabs/extract_frames.py index 702be26149..e41e28ac4f 100644 --- a/deeplabcut/gui/tabs/extract_frames.py +++ b/deeplabcut/gui/tabs/extract_frames.py @@ -91,7 +91,11 @@ def _set_page(self): self._generate_layout_attributes(self.layout_attributes) self.main_layout.addLayout(self.layout_attributes) - self.main_layout.addWidget(_create_label_widget("Optional: frame extraction from a video subset", "font:bold")) + self.main_layout.addWidget( + _create_label_widget( + "Optional: frame extraction from a video subset", "font:bold" + ) + ) self.video_selection_widget = VideoSelectionWidget(self.root, self) self.main_layout.addWidget(self.video_selection_widget) diff --git a/deeplabcut/gui/tracklet_toolbox.py b/deeplabcut/gui/tracklet_toolbox.py index 34b2785cb0..388ffbeb42 100644 --- a/deeplabcut/gui/tracklet_toolbox.py +++ b/deeplabcut/gui/tracklet_toolbox.py @@ -358,7 +358,7 @@ def _prepare_canvas(self, manager, fig): img = self.video.read_frame() self.im = self.ax1.imshow(img) - self.scat = self.ax1.scatter([], [], s=self.dotsize**2, picker=True) + self.scat = self.ax1.scatter([], [], s=self.dotsize ** 2, picker=True) self.scat.set_offsets(manager.xy[:, 0]) self.scat.set_color(self.colors) self.trails = sum( @@ -807,7 +807,7 @@ def on_change(self, val): def update_dotsize(self, val): self.dotsize = val - self.scat.set_sizes([self.dotsize**2]) + self.scat.set_sizes([self.dotsize ** 2]) @staticmethod def calc_distance(x1, y1, x2, y2): diff --git a/deeplabcut/gui/window.py b/deeplabcut/gui/window.py index 71f291c73e..fa458d79d4 100644 --- a/deeplabcut/gui/window.py +++ b/deeplabcut/gui/window.py @@ -407,13 +407,17 @@ def _update_project_state(self, config, loaded): def _ask_for_help(self): dlg = QMessageBox(self) dlg.setWindowTitle("Ask for help") - dlg.setText('''Ask our community for help on the forum!''') + dlg.setText( + """Ask our community for help on the forum!""" + ) _ = dlg.exec() def _learn_dlc(self): dlg = QMessageBox(self) dlg.setWindowTitle("Learn DLC") - dlg.setText('''Learn DLC with our docs and how-to guides!''') + dlg.setText( + """Learn DLC with our docs and how-to guides!""" + ) _ = dlg.exec() def _create_project(self): diff --git a/deeplabcut/modelzoo/api/spatiotemporal_adapt.py b/deeplabcut/modelzoo/api/spatiotemporal_adapt.py index ab779585f7..bef6146ba2 100644 --- a/deeplabcut/modelzoo/api/spatiotemporal_adapt.py +++ b/deeplabcut/modelzoo/api/spatiotemporal_adapt.py @@ -119,7 +119,10 @@ def before_adapt_inference(self, make_video=False, **kwargs): customized_test_config=self.customized_pose_config, ) if kwargs.pop("plot_trajectories", True): - _plot_trajectories(datafiles[0]) + if len(datafiles) == 0: + print("No data files found for plotting trajectory") + else: + _plot_trajectories(datafiles[0]) if make_video: deeplabcut.create_labeled_video( diff --git a/deeplabcut/modelzoo/api/superanimal_inference.py b/deeplabcut/modelzoo/api/superanimal_inference.py index dc14533329..0d6ebe97fa 100644 --- a/deeplabcut/modelzoo/api/superanimal_inference.py +++ b/deeplabcut/modelzoo/api/superanimal_inference.py @@ -341,7 +341,7 @@ def video_inference( print("Loading ", video) vid = VideoWriter(video) if len(scale_list) == 0: - # spatial pyramid can still be useful for reducing jittering and quantization error + # spatial pyramid can still be useful for reducing jittering and quantization error scale_list = [vid.height - 50, vid.height, vid.height + 50] if robust_nframes: nframes = vid.get_n_frames(robust=True) diff --git a/deeplabcut/pose_estimation_tensorflow/core/train_multianimal.py b/deeplabcut/pose_estimation_tensorflow/core/train_multianimal.py index 08ea4ce7fc..79f96c539a 100644 --- a/deeplabcut/pose_estimation_tensorflow/core/train_multianimal.py +++ b/deeplabcut/pose_estimation_tensorflow/core/train_multianimal.py @@ -46,7 +46,7 @@ def train( traintime_resize=False, video_path="", superanimal=None, - trim_ends = None # trim the both ends of the video for video adaptation + remove_head=False, ): # in case there was already a graph tf.compat.v1.reset_default_graph() @@ -93,6 +93,7 @@ def train( cfg["pairwise_predict"] = True dataset = PoseDatasetFactory.create(cfg) + batch_spec = get_batch_spec(cfg) batch, enqueue_op, placeholders = setup_preloading(batch_spec) @@ -107,6 +108,7 @@ def train( if init_weights != "": cfg["init_weights"] = init_weights cfg["resume_weights_only"] = True + print("replacing default init weights with: ", init_weights) stem = Path(cfg["init_weights"]).stem if "snapshot" in stem and keepdeconvweights: @@ -117,6 +119,14 @@ def train( else: start_iter = int(stem.split("-")[1]) + if remove_head: + # removing the decoding layer from the checkpoint + temp = [] + for variable in variables_to_restore: + if "pose" not in variable.name: + temp.append(variable) + variables_to_restore = temp + else: print("Loading ImageNet-pretrained", net_type) # loading backbone from ResNet, MobileNet etc. diff --git a/deeplabcut/pose_estimation_tensorflow/datasets/pose_deterministic.py b/deeplabcut/pose_estimation_tensorflow/datasets/pose_deterministic.py index 146a288333..681150e919 100644 --- a/deeplabcut/pose_estimation_tensorflow/datasets/pose_deterministic.py +++ b/deeplabcut/pose_estimation_tensorflow/datasets/pose_deterministic.py @@ -228,7 +228,7 @@ def make_batch(self, data_item, scale, mirror): def compute_target_part_scoremap(self, joint_id, coords, data_item, size, scale): dist_thresh = self.cfg["pos_dist_thresh"] * scale - dist_thresh_sq = dist_thresh**2 + dist_thresh_sq = dist_thresh ** 2 num_joints = self.cfg["num_joints"] scmap = np.zeros(np.concatenate([size, np.array([num_joints])])) locref_size = np.concatenate([size, np.array([num_joints * 2])]) @@ -260,7 +260,7 @@ def compute_target_part_scoremap(self, joint_id, coords, data_item, size, scale) pt_x = i * self.stride + self.half_stride dx = j_x - pt_x dy = j_y - pt_y - dist = dx**2 + dy**2 + dist = dx ** 2 + dy ** 2 # print(la.norm(diff)) if dist <= dist_thresh_sq: scmap[j, i, j_id] = 1 diff --git a/deeplabcut/pose_estimation_tensorflow/datasets/pose_imgaug.py b/deeplabcut/pose_estimation_tensorflow/datasets/pose_imgaug.py index db60ec8090..11118e02d9 100644 --- a/deeplabcut/pose_estimation_tensorflow/datasets/pose_imgaug.py +++ b/deeplabcut/pose_estimation_tensorflow/datasets/pose_imgaug.py @@ -487,7 +487,7 @@ def gaussian_scmap(self, joint_id, coords, data_item, size, scale): width = size[1] height = size[0] dist_thresh = float((width + height) / 6) - dist_thresh_sq = dist_thresh**2 + dist_thresh_sq = dist_thresh ** 2 std = dist_thresh / 4 # Grid of coordinates @@ -503,7 +503,7 @@ def gaussian_scmap(self, joint_id, coords, data_item, size, scale): map_j = grid.copy() # Distance between the joint point and each coordinate dist = np.linalg.norm(grid - (j_y, j_x), axis=2) ** 2 - scmap_j = np.exp(-dist / (2 * (std**2))) + scmap_j = np.exp(-dist / (2 * (std ** 2))) scmap[..., j_id] = scmap_j locref_mask[dist <= dist_thresh_sq, j_id * 2 + 0] = 1 locref_mask[dist <= dist_thresh_sq, j_id * 2 + 1] = 1 @@ -528,7 +528,7 @@ def compute_target_part_scoremap_numpy( self, joint_id, coords, data_item, size, scale ): dist_thresh = float(self.cfg["pos_dist_thresh"] * scale) - dist_thresh_sq = dist_thresh**2 + dist_thresh_sq = dist_thresh ** 2 num_joints = self.cfg["num_joints"] scmap = np.zeros(np.concatenate([size, np.array([num_joints])])) @@ -555,7 +555,7 @@ def compute_target_part_scoremap_numpy( y = grid.copy()[:, :, 0] dx = j_x - x * self.stride - self.half_stride dy = j_y - y * self.stride - self.half_stride - dist = dx**2 + dy**2 + dist = dx ** 2 + dy ** 2 mask1 = dist <= dist_thresh_sq mask2 = (x >= min_x) & (x <= max_x) mask3 = (y >= min_y) & (y <= max_y) diff --git a/deeplabcut/pose_estimation_tensorflow/datasets/pose_multianimal_imgaug.py b/deeplabcut/pose_estimation_tensorflow/datasets/pose_multianimal_imgaug.py index 8659a05314..5a973a39f6 100644 --- a/deeplabcut/pose_estimation_tensorflow/datasets/pose_multianimal_imgaug.py +++ b/deeplabcut/pose_estimation_tensorflow/datasets/pose_multianimal_imgaug.py @@ -131,12 +131,14 @@ def load_dataset(self): self.has_gt = has_gt return data - def _load_pseudo_data_from_h5(self, cfg, threshold=0.5, mask_kpts_below_thresh=False): + def _load_pseudo_data_from_h5( + self, cfg, threshold=0.5, mask_kpts_below_thresh=False + ): gt_file = cfg["pseudo_label"] assert os.path.exists(gt_file) path_ = Path(gt_file) print("Using gt file:", path_.name) - num_kpts = len(cfg['all_joints_names']) + num_kpts = len(cfg["all_joints_names"]) df = pd.read_hdf(gt_file) video_name = path_.name.split("DLC")[0] video_root = str(path_.parents[0] / video_name) @@ -353,11 +355,13 @@ def get_batch_from_video(self): batch_joints = [] joint_ids = [] data_items = [] - trim_ends = self.cfg.get('trim_ends', None) + trim_ends = self.cfg.get("trim_ends", None) if trim_ends is None: trim_ends = 0 # because of the existence of threshold, sampling population is adjusted to len(self.data) - img_idx = np.random.choice(len(self.data) - trim_ends *2, size=self.batch_size, replace=True) + img_idx = np.random.choice( + len(self.data) - trim_ends * 2, size=self.batch_size, replace=True + ) for i in range(self.batch_size): index = img_idx[i] offset = trim_ends @@ -614,7 +618,7 @@ def compute_target_part_scoremap_numpy( locref_size = *size, num_joints * 2 locref_map = np.zeros(locref_size) locref_scale = 1.0 / self.cfg["locref_stdev"] - dist_thresh_sq = dist_thresh**2 + dist_thresh_sq = dist_thresh ** 2 partaffinityfield_shape = *size, self.cfg["num_limbs"] * 2 partaffinityfield_map = np.zeros(partaffinityfield_shape) @@ -640,7 +644,7 @@ def compute_target_part_scoremap_numpy( dx_ = dx * locref_scale dy = coords[:, 1] - yy * stride - half_stride dy_ = dy * locref_scale - dist = dx**2 + dy**2 + dist = dx ** 2 + dy ** 2 mask1 = dist <= dist_thresh_sq mask2 = (xx >= mins[:, 0]) & (xx <= maxs[:, 0]) mask3 = (yy >= mins[:, 1]) & (yy <= maxs[:, 1]) @@ -747,7 +751,7 @@ def gaussian_scmap(self, joint_id, coords, data_item, size, scale): locref_map = np.zeros(locref_size) locref_scale = 1.0 / self.cfg["locref_stdev"] - dist_thresh_sq = dist_thresh**2 + dist_thresh_sq = dist_thresh ** 2 partaffinityfield_shape = np.concatenate( [size, np.array([self.cfg["num_limbs"] * 2])] @@ -779,7 +783,7 @@ def gaussian_scmap(self, joint_id, coords, data_item, size, scale): map_j = grid.copy() # Distance between the joint point and each coordinate dist = np.linalg.norm(grid - (j_y, j_x), axis=2) ** 2 - scmap_j = np.exp(-dist / (2 * (std**2))) + scmap_j = np.exp(-dist / (2 * (std ** 2))) scmap[..., j_id] = scmap_j locref_mask[dist <= dist_thresh_sq, j_id * 2 + 0] = 1 locref_mask[dist <= dist_thresh_sq, j_id * 2 + 1] = 1 diff --git a/deeplabcut/pose_estimation_tensorflow/datasets/pose_tensorpack.py b/deeplabcut/pose_estimation_tensorflow/datasets/pose_tensorpack.py index d25e98a6a4..5a1e591fd2 100644 --- a/deeplabcut/pose_estimation_tensorflow/datasets/pose_tensorpack.py +++ b/deeplabcut/pose_estimation_tensorflow/datasets/pose_tensorpack.py @@ -350,7 +350,7 @@ def compute_target_part_scoremap(self, components): locref_map = np.zeros(locref_size) locref_scale = 1.0 / self.cfg["locref_stdev"] - dist_thresh_sq = dist_thresh**2 + dist_thresh_sq = dist_thresh ** 2 width = size[1] height = size[0] @@ -375,7 +375,7 @@ def compute_target_part_scoremap(self, components): pt_x = i * stride + half_stride dx = j_x - pt_x dy = j_y - pt_y - dist = dx**2 + dy**2 + dist = dx ** 2 + dy ** 2 # print(la.norm(diff)) if dist <= dist_thresh_sq: scmap[j, i, j_id] = 1 diff --git a/deeplabcut/pose_estimation_tensorflow/lib/inferenceutils.py b/deeplabcut/pose_estimation_tensorflow/lib/inferenceutils.py index 4565a576bd..8a7dc4a690 100644 --- a/deeplabcut/pose_estimation_tensorflow/lib/inferenceutils.py +++ b/deeplabcut/pose_estimation_tensorflow/lib/inferenceutils.py @@ -359,7 +359,7 @@ def calc_link_probability(self, link): ind = _conv_square_to_condensed_indices(i, j, self.n_multibodyparts) mu = self._kde.mean[ind] sigma = self._kde.covariance[ind, ind] - z = (link.length**2 - mu) / sigma + z = (link.length ** 2 - mu) / sigma return 2 * (1 - 0.5 * (1 + erf(abs(z) / sqrt(2)))) @staticmethod @@ -809,13 +809,13 @@ def wrapped(i): if unique is not None: self.unique[i] = unique pbar.update() - + def from_pickle(self, pickle_path): with open(pickle_path, "rb") as file: data = pickle.load(file) self.unique = data.pop("single", {}) self.assemblies = data - + @staticmethod def parse_metadata(data): params = dict() diff --git a/deeplabcut/pose_estimation_tensorflow/nnets/utils.py b/deeplabcut/pose_estimation_tensorflow/nnets/utils.py index bb080952df..8afd4c4552 100644 --- a/deeplabcut/pose_estimation_tensorflow/nnets/utils.py +++ b/deeplabcut/pose_estimation_tensorflow/nnets/utils.py @@ -86,8 +86,8 @@ def get_batch_spec(cfg): def make_2d_gaussian_kernel(sigma, size): sigma = tf.convert_to_tensor(sigma, dtype=tf.float32) k = tf.range(-size // 2 + 1, size // 2 + 1) - k = tf.cast(k**2, sigma.dtype) - k = tf.nn.softmax(-k / (2 * (sigma**2))) + k = tf.cast(k ** 2, sigma.dtype) + k = tf.nn.softmax(-k / (2 * (sigma ** 2))) return tf.einsum("i,j->ij", k, k) diff --git a/deeplabcut/pose_estimation_tensorflow/predict_supermodel.py b/deeplabcut/pose_estimation_tensorflow/predict_supermodel.py index e282ee3362..48092ea82e 100644 --- a/deeplabcut/pose_estimation_tensorflow/predict_supermodel.py +++ b/deeplabcut/pose_estimation_tensorflow/predict_supermodel.py @@ -20,9 +20,9 @@ def video_inference_superanimal( video_adapt=False, plot_trajectories=True, pcutoff=0.1, - adapt_iterations = 1000, - pseudo_threshold = 0.1, - trim_ends = None + adapt_iterations=1000, + pseudo_threshold=0.1, + trim_ends=None, ): """ Makes prediction based on a super animal model. Note right now we only support single animal video inference @@ -63,7 +63,7 @@ def video_inference_superanimal( trim_ends: int, optional: In cases where the beginning and ending of the videos have very messy background that impacts predictions of the model, we trim those from adaptation training - + Given a list of scales for spatial pyramid, i.e. [600, 700] scale_list = range(600,800,100) @@ -99,16 +99,17 @@ def video_inference_superanimal( scale_list=scale_list, ) if not video_adapt: - adapter.before_adapt_inference(make_video=True, - pcutoff=pcutoff, - plot_trajectories = plot_trajectories) + adapter.before_adapt_inference( + make_video=True, pcutoff=pcutoff, plot_trajectories=plot_trajectories + ) else: adapter.before_adapt_inference(make_video=False) - adapter.adaptation_training(adapt_iterations = adapt_iterations, - pseudo_threshold = pseudo_threshold, - trim_ends = trim_ends) + adapter.adaptation_training( + adapt_iterations=adapt_iterations, + pseudo_threshold=pseudo_threshold, + trim_ends=trim_ends, + ) adapter.after_adapt_inference( pcutoff=pcutoff, plot_trajectories=plot_trajectories, ) - diff --git a/deeplabcut/pose_estimation_tensorflow/predict_videos.py b/deeplabcut/pose_estimation_tensorflow/predict_videos.py index ff42a5b7f8..3a994388fa 100644 --- a/deeplabcut/pose_estimation_tensorflow/predict_videos.py +++ b/deeplabcut/pose_estimation_tensorflow/predict_videos.py @@ -1529,7 +1529,9 @@ def _convert_detections_to_tracklets( assemblies = assembly_builder.assemblies.get(i) if assemblies is None: continue - animals = np.stack([assembly_builder.data[:, :3] for assembly_builder in assemblies]) + animals = np.stack( + [assembly_builder.data[:, :3] for assembly_builder in assemblies] + ) if track_method == "box": xy = trackingutils.calc_bboxes_from_keypoints( animals, inference_cfg.get("boundingboxslack", 0) @@ -1813,7 +1815,9 @@ def convert_detections2tracklets( assemblies_filename = dataname.split(".h5")[0] + "_assemblies.pickle" if not os.path.exists(assemblies_filename) or overwrite: if calibrate: - trainingsetfolder = auxiliaryfunctions.get_training_set_folder(cfg) + trainingsetfolder = auxiliaryfunctions.get_training_set_folder( + cfg + ) train_data_file = os.path.join( cfg["project_path"], str(trainingsetfolder), @@ -1857,7 +1861,9 @@ def convert_detections2tracklets( assemblies = assembly_builder.assemblies.get(index) if assemblies is None: continue - animals = np.stack([assembly_builder.data for assembly_builder in assemblies]) + animals = np.stack( + [assembly_builder.data for assembly_builder in assemblies] + ) if not identity_only: if track_method == "box": xy = trackingutils.calc_bboxes_from_keypoints( diff --git a/deeplabcut/pose_estimation_tensorflow/training.py b/deeplabcut/pose_estimation_tensorflow/training.py index 8ee2137bbc..7ac99815ca 100644 --- a/deeplabcut/pose_estimation_tensorflow/training.py +++ b/deeplabcut/pose_estimation_tensorflow/training.py @@ -63,87 +63,96 @@ def train_network( autotune=False, keepdeconvweights=True, modelprefix="", + superanimal_name="", + superanimal_transfer_learning=False, ): """Trains the network with the labels in the training dataset. - Parameters - ---------- - config : string - Full path of the config.yaml file as a string. - - shuffle: int, optional, default=1 - Integer value specifying the shuffle index to select for training. - - trainingsetindex: int, optional, default=0 - Integer specifying which TrainingsetFraction to use. - Note that TrainingFraction is a list in config.yaml. - - max_snapshots_to_keep: int or None - Sets how many snapshots are kept, i.e. states of the trained network. Every - saving iteration many times a snapshot is stored, however only the last - ``max_snapshots_to_keep`` many are kept! If you change this to None, then all - are kept. - See: https://github.com/DeepLabCut/DeepLabCut/issues/8#issuecomment-387404835 - - displayiters: optional, default=None - This variable is actually set in ``pose_config.yaml``. However, you can - overwrite it with this hack. Don't use this regularly, just if you are too lazy - to dig out the ``pose_config.yaml`` file for the corresponding project. If - ``None``, the value from there is used, otherwise it is overwritten! - - saveiters: optional, default=None - This variable is actually set in ``pose_config.yaml``. However, you can - overwrite it with this hack. Don't use this regularly, just if you are too lazy - to dig out the ``pose_config.yaml`` file for the corresponding project. - If ``None``, the value from there is used, otherwise it is overwritten! - - maxiters: optional, default=None - This variable is actually set in ``pose_config.yaml``. However, you can - overwrite it with this hack. Don't use this regularly, just if you are too lazy - to dig out the ``pose_config.yaml`` file for the corresponding project. - If ``None``, the value from there is used, otherwise it is overwritten! - - allow_growth: bool, optional, default=True. - For some smaller GPUs the memory issues happen. If ``True``, the memory - allocator does not pre-allocate the entire specified GPU memory region, instead - starting small and growing as needed. - See issue: https://forum.image.sc/t/how-to-stop-running-out-of-vram/30551/2 - - gputouse: optional, default=None - Natural number indicating the number of your GPU (see number in nvidia-smi). - If you do not have a GPU put None. - See: https://nvidia.custhelp.com/app/answers/detail/a_id/3751/~/useful-nvidia-smi-queries - - autotune: bool, optional, default=False - Property of TensorFlow, somehow faster if ``False`` - (as Eldar found out, see https://github.com/tensorflow/tensorflow/issues/13317). - - keepdeconvweights: bool, optional, default=True - Also restores the weights of the deconvolution layers (and the backbone) when - training from a snapshot. Note that if you change the number of bodyparts, you - need to set this to false for re-training. - - modelprefix: str, optional, default="" - Directory containing the deeplabcut models to use when evaluating the network. - By default, the models are assumed to exist in the project folder. - - Returns - ------- - None - - Examples - -------- - To train the network for first shuffle of the training dataset - - >>> deeplabcut.train_network('/analysis/project/reaching-task/config.yaml') - - To train the network for second shuffle of the training dataset - - >>> deeplabcut.train_network( - '/analysis/project/reaching-task/config.yaml', - shuffle=2, - keepdeconvweights=True, - ) + Parameters + ---------- + config : string + Full path of the config.yaml file as a string. + + shuffle: int, optional, default=1 + Integer value specifying the shuffle index to select for training. + + trainingsetindex: int, optional, default=0 + Integer specifying which TrainingsetFraction to use. + Note that TrainingFraction is a list in config.yaml. + + max_snapshots_to_keep: int or None + Sets how many snapshots are kept, i.e. states of the trained network. Every + saving iteration many times a snapshot is stored, however only the last + ``max_snapshots_to_keep`` many are kept! If you change this to None, then all + are kept. + See: https://github.com/DeepLabCut/DeepLabCut/issues/8#issuecomment-387404835 + + displayiters: optional, default=None + This variable is actually set in ``pose_config.yaml``. However, you can + overwrite it with this hack. Don't use this regularly, just if you are too lazy + to dig out the ``pose_config.yaml`` file for the corresponding project. If + ``None``, the value from there is used, otherwise it is overwritten! + + saveiters: optional, default=None + This variable is actually set in ``pose_config.yaml``. However, you can + overwrite it with this hack. Don't use this regularly, just if you are too lazy + to dig out the ``pose_config.yaml`` file for the corresponding project. + If ``None``, the value from there is used, otherwise it is overwritten! + + maxiters: optional, default=None + This variable is actually set in ``pose_config.yaml``. However, you can + overwrite it with this hack. Don't use this regularly, just if you are too lazy + to dig out the ``pose_config.yaml`` file for the corresponding project. + If ``None``, the value from there is used, otherwise it is overwritten! + + allow_growth: bool, optional, default=True. + For some smaller GPUs the memory issues happen. If ``True``, the memory + allocator does not pre-allocate the entire specified GPU memory region, instead + starting small and growing as needed. + See issue: https://forum.image.sc/t/how-to-stop-running-out-of-vram/30551/2 + + gputouse: optional, default=None + Natural number indicating the number of your GPU (see number in nvidia-smi). + If you do not have a GPU put None. + See: https://nvidia.custhelp.com/app/answers/detail/a_id/3751/~/useful-nvidia-smi-queries + + autotune: bool, optional, default=False + Property of TensorFlow, somehow faster if ``False`` + (as Eldar found out, see https://github.com/tensorflow/tensorflow/issues/13317). + + keepdeconvweights: bool, optional, default=True + Also restores the weights of the deconvolution layers (and the backbone) when + training from a snapshot. Note that if you change the number of bodyparts, you + need to set this to false for re-training. + + modelprefix: str, optional, default="" + Directory containing the deeplabcut models to use when evaluating the network. + By default, the models are assumed to exist in the project folder. + + superanimal_name: str, optional, default ="" + Specified if transfer learning with superanimal is desired + + superanimal_transfer_learning: bool, optional, default = False. + If set true, the training is transfer learning (new decoding layer). If set false, + and superanimal_name is True, then the training is fine-tuning (reusing the decoding layer) + + Returns + ------- + None + + Examples + -------- + To train the network for first shuffle of the training dataset + + >>> deeplabcut.train_network('/analysis/project/reaching-task/config.yaml') + + To train the network for second shuffle of the training dataset + + >>> deeplabcut.train_network( + '/analysis/project/reaching-task/config.yaml', + shuffle=2, + keepdeconvweights=True, + ) """ if allow_growth: os.environ["TF_FORCE_GPU_ALLOW_GROWTH"] = "true" @@ -190,7 +199,59 @@ def train_network( os.environ["CUDA_VISIBLE_DEVICES"] = str(gputouse) try: cfg_dlc = auxiliaryfunctions.read_plainconfig(poseconfigfile) - if "multi-animal" in cfg_dlc["dataset_type"]: + + if superanimal_name != "": + from deeplabcut.modelzoo.utils import parse_available_supermodels + from dlclibrary.dlcmodelzoo.modelzoo_download import ( + download_huggingface_model, + MODELOPTIONS, + ) + import glob + + dlc_root_path = auxiliaryfunctions.get_deeplabcut_path() + supermodels = parse_available_supermodels() + weight_folder = str( + Path(dlc_root_path) + / "pose_estimation_tensorflow" + / "models" + / "pretrained" + / (superanimal_name + "_weights") + ) + + if superanimal_name in MODELOPTIONS: + if not os.path.exists(weight_folder): + download_huggingface_model(superanimal_name, weight_folder) + else: + print(f"{weight_folder} exists, using the downloaded weights") + else: + print( + f"{superanimal_name} not available. Available ones are: ", + MODELOPTIONS, + ) + + snapshots = glob.glob(os.path.join(weight_folder, "snapshot-*.index")) + init_weights = os.path.abspath(snapshots[0]).replace(".index", "") + + from deeplabcut.pose_estimation_tensorflow.core.train_multianimal import ( + train, + ) + + print("Selecting multi-animal trainer") + train( + str(poseconfigfile), + displayiters, + saveiters, + maxiters, + max_to_keep=max_snapshots_to_keep, + keepdeconvweights=keepdeconvweights, + allow_growth=allow_growth, + init_weights=init_weights, + remove_head=True + if superanimal_name != "" and superanimal_transfer_learning + else False, + ) # pass on path and file name for pose_cfg.yaml! + + elif "multi-animal" in cfg_dlc["dataset_type"]: from deeplabcut.pose_estimation_tensorflow.core.train_multianimal import ( train, ) diff --git a/deeplabcut/pose_estimation_tensorflow/util/visualize.py b/deeplabcut/pose_estimation_tensorflow/util/visualize.py index 12fb8cd7b7..29fd039fc2 100644 --- a/deeplabcut/pose_estimation_tensorflow/util/visualize.py +++ b/deeplabcut/pose_estimation_tensorflow/util/visualize.py @@ -31,7 +31,7 @@ def _npcircle(image, cx, cy, radius, color, transparency=0.0): cx = int(cx) cy = int(cy) y, x = np.ogrid[-radius:radius, -radius:radius] - index = x**2 + y**2 <= radius**2 + index = x ** 2 + y ** 2 <= radius ** 2 image[cy - radius : cy + radius, cx - radius : cx + radius][index] = ( image[cy - radius : cy + radius, cx - radius : cx + radius][index].astype( "float32" diff --git a/deeplabcut/pose_estimation_tensorflow/visualizemaps.py b/deeplabcut/pose_estimation_tensorflow/visualizemaps.py index 99ef0dc549..5c417e5ef2 100644 --- a/deeplabcut/pose_estimation_tensorflow/visualizemaps.py +++ b/deeplabcut/pose_estimation_tensorflow/visualizemaps.py @@ -339,7 +339,7 @@ def visualize_paf(image, paf, step=5, colors=None): V = paf[:, :, n, 1] X, Y = np.meshgrid(np.arange(U.shape[1]), np.arange(U.shape[0])) M = np.zeros(U.shape, dtype=bool) - M[U**2 + V**2 < 0.5 * 0.5**2] = True + M[U ** 2 + V ** 2 < 0.5 * 0.5 ** 2] = True U = np.ma.masked_array(U, mask=M) V = np.ma.masked_array(V, mask=M) ax.quiver( diff --git a/deeplabcut/pose_tracking_pytorch/model/backbones/vit_pytorch.py b/deeplabcut/pose_tracking_pytorch/model/backbones/vit_pytorch.py index c036503dc7..ef68ffd81f 100644 --- a/deeplabcut/pose_tracking_pytorch/model/backbones/vit_pytorch.py +++ b/deeplabcut/pose_tracking_pytorch/model/backbones/vit_pytorch.py @@ -111,7 +111,7 @@ def __init__( self.num_heads = num_heads head_dim = dim // num_heads # NOTE scale factor was wrong in my original version, can set manually to be compat with prev weights - self.scale = qk_scale or head_dim**-0.5 + self.scale = qk_scale or head_dim ** -0.5 self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias) self.attn_drop = nn.Dropout(attn_drop) diff --git a/deeplabcut/pose_tracking_pytorch/solver/cosine_lr.py b/deeplabcut/pose_tracking_pytorch/solver/cosine_lr.py index f61d248ea0..7e7aeba855 100644 --- a/deeplabcut/pose_tracking_pytorch/solver/cosine_lr.py +++ b/deeplabcut/pose_tracking_pytorch/solver/cosine_lr.py @@ -96,14 +96,14 @@ def _get_lr(self, t): i = math.floor( math.log(1 - t / self.t_initial * (1 - self.t_mul), self.t_mul) ) - t_i = self.t_mul**i * self.t_initial - t_curr = t - (1 - self.t_mul**i) / (1 - self.t_mul) * self.t_initial + t_i = self.t_mul ** i * self.t_initial + t_curr = t - (1 - self.t_mul ** i) / (1 - self.t_mul) * self.t_initial else: i = t // self.t_initial t_i = self.t_initial t_curr = t - (self.t_initial * i) - gamma = self.decay_rate**i + gamma = self.decay_rate ** i lr_min = self.lr_min * gamma lr_max_values = [v * gamma for v in self.base_values] @@ -139,6 +139,6 @@ def get_cycle_length(self, cycles=0): else: return int( math.floor( - -self.t_initial * (self.t_mul**cycles - 1) / (1 - self.t_mul) + -self.t_initial * (self.t_mul ** cycles - 1) / (1 - self.t_mul) ) ) diff --git a/deeplabcut/pose_tracking_pytorch/tracking_utils/reranking.py b/deeplabcut/pose_tracking_pytorch/tracking_utils/reranking.py index 957be31efe..e290a494ee 100644 --- a/deeplabcut/pose_tracking_pytorch/tracking_utils/reranking.py +++ b/deeplabcut/pose_tracking_pytorch/tracking_utils/reranking.py @@ -76,7 +76,9 @@ def re_ranking( k_reciprocal_expansion_index = np.unique(k_reciprocal_expansion_index) weight = np.exp(-original_dist[i, k_reciprocal_expansion_index]) V[i, k_reciprocal_expansion_index] = weight / np.sum(weight) - original_dist = original_dist[:query_num,] + original_dist = original_dist[ + :query_num, + ] if k2 != 1: V_qe = np.zeros_like(V, dtype=np.float16) for i in range(all_num): diff --git a/deeplabcut/refine_training_dataset/outlier_frames.py b/deeplabcut/refine_training_dataset/outlier_frames.py index f3cd4be621..687ce93c36 100644 --- a/deeplabcut/refine_training_dataset/outlier_frames.py +++ b/deeplabcut/refine_training_dataset/outlier_frames.py @@ -407,7 +407,7 @@ def extract_outlier_frames( temp_dt = df_temp.diff(axis=0) ** 2 temp_dt.drop("likelihood", axis=1, level="coords", inplace=True) sum_ = temp_dt.groupby(level="bodyparts", axis=1).sum() - ind = df_temp.index[(sum_ > epsilon**2).any(axis=1)].tolist() + ind = df_temp.index[(sum_ > epsilon ** 2).any(axis=1)].tolist() Indices.extend(ind) elif outlieralgorithm == "fitting": d, o = compute_deviations( @@ -993,7 +993,7 @@ def PlottingSingleFrame( plt.scatter( df_x[ind, index], df_y[ind, index], - s=dotsize**2, + s=dotsize ** 2, color=colors(map2bp[i]), alpha=alphavalue, ) @@ -1065,7 +1065,7 @@ def PlottingSingleFramecv2( plt.scatter( df_x[ind, index], df_y[ind, index], - s=dotsize**2, + s=dotsize ** 2, color=colors(map2bp[i]), alpha=alphavalue, ) diff --git a/deeplabcut/refine_training_dataset/stitch.py b/deeplabcut/refine_training_dataset/stitch.py index 54ce652a98..d760f1efd1 100644 --- a/deeplabcut/refine_training_dataset/stitch.py +++ b/deeplabcut/refine_training_dataset/stitch.py @@ -222,13 +222,13 @@ def calc_velocity(self, where="head", norm=True): else: raise ValueError(f"Unknown where={where}") if norm: - return np.sqrt(np.sum(vel**2, axis=1)).mean() + return np.sqrt(np.sum(vel ** 2, axis=1)).mean() return vel.mean(axis=0) @property def maximal_velocity(self): vel = np.diff(self.centroid, axis=0) / np.diff(self.inds)[:, np.newaxis] - return np.sqrt(np.max(np.sum(vel**2, axis=1))) + return np.sqrt(np.max(np.sum(vel ** 2, axis=1))) def calc_rate_of_turn(self, where="head"): """ @@ -267,7 +267,7 @@ def distance_to(self, other_tracklet): self.centroid[np.isin(self.inds, other_tracklet.inds)] - other_tracklet.centroid[np.isin(other_tracklet.inds, self.inds)] ) - return np.sqrt(np.sum(dist**2, axis=1)).mean() + return np.sqrt(np.sum(dist ** 2, axis=1)).mean() elif self < other_tracklet: return np.sqrt( np.sum((self.centroid[-1] - other_tracklet.centroid[0]) ** 2) @@ -300,7 +300,7 @@ def motion_affinity_with(self, other_tracklet): d2 = self.centroid[0] - time_gap * self.calc_velocity("tail", False) delta1 = self.centroid[0] - d1 delta2 = other_tracklet.centroid[-1] - d2 - return (np.sqrt(np.sum(delta1**2)) + np.sqrt(np.sum(delta2**2))) / 2 + return (np.sqrt(np.sum(delta1 ** 2)) + np.sqrt(np.sum(delta2 ** 2))) / 2 return 0 def time_gap_to(self, other_tracklet): @@ -417,7 +417,7 @@ def estimate_rank(self, tol): # omega = 0.56 * beta ** 3 - 0.95 * beta ** 2 + 1.82 * beta + 1.43 _, s, _ = sli.svd(mat, min(10, min(mat.shape))) # return np.argmin(s > omega * np.median(s)) - eigen = s**2 + eigen = s ** 2 diff = np.abs(np.diff(eigen / eigen[0])) return np.argmin(diff > tol) @@ -873,7 +873,7 @@ def concatenate_data(self): temp = np.full((self.n_frames, flat_data.shape[1]), np.nan) temp[track.inds - self._first_frame] = flat_data data.append(temp) - + # If there isn't a track for each animal, fill in the dataframe with NaNs missing_tracks = self.n_tracks - len(self.tracks) if missing_tracks > 0: diff --git a/deeplabcut/utils/make_labeled_video.py b/deeplabcut/utils/make_labeled_video.py index 514fe58d05..238fd4e3ef 100644 --- a/deeplabcut/utils/make_labeled_video.py +++ b/deeplabcut/utils/make_labeled_video.py @@ -332,14 +332,14 @@ def CreateVideoSlow( ax.scatter( df_x[ind][max(0, index - trailpoints) : index], df_y[ind][max(0, index - trailpoints) : index], - s=dotsize**2, + s=dotsize ** 2, color=color, alpha=alphavalue * 0.75, ) ax.scatter( df_x[ind, index], df_y[ind, index], - s=dotsize**2, + s=dotsize ** 2, color=color, alpha=alphavalue, ) @@ -967,7 +967,7 @@ def create_video_with_keypoints_only( plt.switch_backend("agg") fig = plt.figure(frameon=False, figsize=(nx / dpi, ny / dpi)) ax = fig.add_subplot(111) - scat = ax.scatter([], [], s=dotsize**2, alpha=alpha) + scat = ax.scatter([], [], s=dotsize ** 2, alpha=alpha) coords = xyp[0, :, :2] coords[xyp[0, :, 2] < pcutoff] = np.nan scat.set_offsets(coords) diff --git a/docs/ModelZoo.md b/docs/ModelZoo.md index 89aa2a4514..0ef8864d1e 100644 --- a/docs/ModelZoo.md +++ b/docs/ModelZoo.md @@ -47,14 +47,38 @@ Via DeepLabCut Model Zoo, we aim to provide plug and play models that do not nee pip install deeplabcut[tf,modelzoo] ``` +#### Practical example: Using SuperAnimal models for inference without training. +In the `deeplabcut.video_inference_superanimal` function, if the output video appears to be jittery, consider setting the `video_adapt` option to __True__. Be aware, that enabling this option might extend the processing time. + ```python video_path = 'demo-video.mp4' superanimal_name = 'superanimal_quadruped' -scale_list = range(200, 600, 50) # image height pixel size range and increment -deeplabcut.video_inference_superanimal([video_path], superanimal_name, scale_list=scale_list) +# The purpose of the scale list is to aggregate predictions from various image sizes. We anticipate the appearance size of the animal in the images to be approximately 400 pixels. +scale_list = range(200, 600, 50) + +deeplabcut.video_inference_superanimal([video_path], superanimal_name, scale_list=scale_list, video_adapt = False) ``` +#### Practical example: Using transfer learning with superanimal weights. +In the `deeplabcut.train_network` function, the `superanimal_transfer_learning` option plays a pivotal role. If it's set to __True__, it uses a new decoding layer and allows you to use superanimal weights in any project, no matter the number of keypoints. However, if it's set to __False__, you are doing fine-tuning. So, make sure your dataset has the right number of keypoints. + Specifically: +* `superquadruped` uses 39 keypoints and, +* `supertopview` uses 27 keypoints + +```python +superanimal_name = "superanimal_topviewmouse" +config_path = os.path.join(os.getcwd(), "openfield-Pranav-2018-10-30", "config.yaml") + +deeplabcut.create_training_dataset(config_path, superanimal_name = superanimal_name) + +deeplabcut.train_network(config_path, + maxiters=10, + superanimal_name = superanimal_name, + superanimal_transfer_learning = True) +``` + + ### To see the list of available models, check out the [Home page](http://modelzoo.deeplabcut.org/). diff --git a/examples/testscript_superanimal_transfer_learning.py b/examples/testscript_superanimal_transfer_learning.py new file mode 100644 index 0000000000..1f36f99074 --- /dev/null +++ b/examples/testscript_superanimal_transfer_learning.py @@ -0,0 +1,31 @@ +# +# DeepLabCut Toolbox (deeplabcut.org) +# © A. & M.W. Mathis Labs +# https://github.com/DeepLabCut/DeepLabCut +# +# Please see AUTHORS for contributors. +# https://github.com/DeepLabCut/DeepLabCut/blob/master/AUTHORS +# +# Licensed under GNU Lesser General Public License v3.0 +# +""" +Test script for super animal adaptation +""" +import deeplabcut +import os + +print(deeplabcut.__file__) +if __name__ == "__main__": + + superanimal_name = "superanimal_topviewmouse" + basepath = os.path.dirname(os.path.realpath(__file__)) + config_path = os.path.join(basepath, "openfield-Pranav-2018-10-30", "config.yaml") + + deeplabcut.create_training_dataset(config_path, superanimal_name=superanimal_name) + + deeplabcut.train_network( + config_path, + maxiters=10, + superanimal_name=superanimal_name, + superanimal_transfer_learning=True, + ) diff --git a/setup.py b/setup.py index fe2fd995bc..428fef611a 100644 --- a/setup.py +++ b/setup.py @@ -55,8 +55,10 @@ ], "openvino": ["openvino-dev==2022.1.0"], "docs": ["numpydoc"], - "tf": ["tensorflow>=2.0,<=2.10"], # Last supported TF version on Windows Native is 2.10 - "apple_mchips": ["tensorflow-macos<2.13.0","tensorflow-metal"], + "tf": [ + "tensorflow>=2.0,<=2.10" + ], # Last supported TF version on Windows Native is 2.10 + "apple_mchips": ["tensorflow-macos<2.13.0", "tensorflow-metal"], "modelzoo": ["huggingface_hub"], }, scripts=["deeplabcut/pose_estimation_tensorflow/models/pretrained/download.sh"], diff --git a/tests/test_dataset_augmentation.py b/tests/test_dataset_augmentation.py index 9dab238b1e..a11dd148d0 100644 --- a/tests/test_dataset_augmentation.py +++ b/tests/test_dataset_augmentation.py @@ -98,10 +98,7 @@ def test_keypoint_horizontal_flip( keypoints=list(map(str, range(12))), symmetric_pairs=pairs, ) - keypoints_aug = aug( - images=[sample_image], - keypoints=[sample_keypoints], - )[ + keypoints_aug = aug(images=[sample_image], keypoints=[sample_keypoints],)[ 1 ][0] temp = keypoints_aug.reshape((3, 12, 2)) diff --git a/tests/test_inferenceutils.py b/tests/test_inferenceutils.py index 2736939f6a..f44aad610f 100644 --- a/tests/test_inferenceutils.py +++ b/tests/test_inferenceutils.py @@ -107,7 +107,7 @@ def test_link(): j1 = inferenceutils.Joint(pos1, conf, idx=idx1) j2 = inferenceutils.Joint(pos2, conf, idx=idx2) link = inferenceutils.Link(j1, j2) - assert link.confidence == conf**2 + assert link.confidence == conf ** 2 assert link.idx == (idx1, idx2) assert link.to_vector() == [*pos1, *pos2] From 780a1db46b36895ad6b66f44b3bdc9f277376410 Mon Sep 17 00:00:00 2001 From: Jessy Lauer <30733203+jeylau@users.noreply.github.com> Date: Tue, 7 Nov 2023 15:46:58 +0100 Subject: [PATCH 09/11] Improvements to the model zoo's GUI tab (#2431) * Add more control widgets --- deeplabcut/gui/tabs/modelzoo.py | 74 +++++++++++++++++++++++++++++++-- 1 file changed, 71 insertions(+), 3 deletions(-) diff --git a/deeplabcut/gui/tabs/modelzoo.py b/deeplabcut/gui/tabs/modelzoo.py index 26bc7ff110..2607fab362 100644 --- a/deeplabcut/gui/tabs/modelzoo.py +++ b/deeplabcut/gui/tabs/modelzoo.py @@ -8,16 +8,21 @@ # # Licensed under GNU Lesser General Public License v3.0 # +import os +from functools import partial + import deeplabcut from PySide6 import QtWidgets from PySide6.QtCore import Qt, Signal, QTimer, QRegularExpression -from PySide6.QtGui import QRegularExpressionValidator +from PySide6.QtGui import QPixmap, QRegularExpressionValidator from deeplabcut.gui.components import ( DefaultTab, VideoSelectionWidget, _create_label_widget, _create_grid_layout, ) +from deeplabcut.gui import BASE_DIR +from deeplabcut.gui.utils import move_to_separate_thread from deeplabcut.modelzoo.utils import parse_available_supermodels @@ -65,17 +70,70 @@ def _set_page(self): validator.validationChanged.connect(self._handle_validation_change) self.scales_line.setValidator(validator) + tooltip_label = QtWidgets.QLabel() + tooltip_label.setPixmap( + QPixmap(os.path.join(BASE_DIR, "assets", "icons", "help2.png")).scaledToWidth(30) + ) + tooltip_label.setToolTip( + "Approximate animal sizes in pixels, for spatial pyramid search. If left blank, defaults to video height +/- 50 pixels", + ) + + self.adapt_checkbox = QtWidgets.QCheckBox("Use video adaptation") + self.adapt_checkbox.setChecked(True) + + pseudo_threshold_label = QtWidgets.QLabel("Pseudo-label confidence threshold") + self.pseudo_threshold_spinbox = QtWidgets.QDoubleSpinBox( + decimals=2, + minimum=0.01, + maximum=1.0, + singleStep=0.05, + value=0.1, + wrapping=True, + ) + self.pseudo_threshold_spinbox.setMaximumWidth(300) + + adapt_iter_label = QtWidgets.QLabel("Number of adaptation iterations") + self.adapt_iter_spinbox = QtWidgets.QSpinBox() + self.adapt_iter_spinbox.setRange(100, 10000) + self.adapt_iter_spinbox.setValue(1000) + self.adapt_iter_spinbox.setSingleStep(100) + self.adapt_iter_spinbox.setGroupSeparatorShown(True) + self.adapt_iter_spinbox.setMaximumWidth(300) + model_settings_layout.addWidget(section_title, 0, 0) model_settings_layout.addWidget(model_combo_text, 1, 0) model_settings_layout.addWidget(self.model_combo, 1, 1) model_settings_layout.addWidget(scales_label, 2, 0) model_settings_layout.addWidget(self.scales_line, 2, 1) + model_settings_layout.addWidget(tooltip_label, 2, 2) + model_settings_layout.addWidget(self.adapt_checkbox, 3, 0) + model_settings_layout.addWidget(pseudo_threshold_label, 4, 0) + model_settings_layout.addWidget(self.pseudo_threshold_spinbox, 4, 1) + model_settings_layout.addWidget(adapt_iter_label, 5, 0) + model_settings_layout.addWidget(self.adapt_iter_spinbox, 5, 1) self.main_layout.addLayout(model_settings_layout) self.run_button = QtWidgets.QPushButton("Run") self.run_button.clicked.connect(self.run_video_adaptation) self.main_layout.addWidget(self.run_button, alignment=Qt.AlignRight) + self.help_button = QtWidgets.QPushButton("Help") + self.help_button.clicked.connect(self.show_help_dialog) + self.main_layout.addWidget(self.help_button, alignment=Qt.AlignLeft) + + def show_help_dialog(self): + dialog = QtWidgets.QDialog(self) + layout = QtWidgets.QVBoxLayout() + label = QtWidgets.QLabel(deeplabcut.video_inference_superanimal.__doc__, self) + scroll = QtWidgets.QScrollArea() + scroll.setVerticalScrollBarPolicy(Qt.ScrollBarAlwaysOn) + scroll.setHorizontalScrollBarPolicy(Qt.ScrollBarAlwaysOff) + scroll.setWidgetResizable(True) + scroll.setWidget(label) + layout.addWidget(scroll) + dialog.setLayout(layout) + dialog.exec_() + def _handle_validation_change(self, state): if state == RegExpValidator.Invalid: color = "red" @@ -109,10 +167,20 @@ def run_video_adaptation(self): supermodel_name = self.model_combo.currentText() videotype = self.video_selection_widget.videotype_widget.currentText() - deeplabcut.video_inference_superanimal( + func = partial( + deeplabcut.video_inference_superanimal, videos, supermodel_name, videotype=videotype, - video_adapt=True, + video_adapt=self.adapt_checkbox.isChecked(), scale_list=scales, + pseudo_threshold=self.pseudo_threshold_spinbox.value(), + adapt_iterations=self.adapt_iter_spinbox.value(), ) + + self.worker, self.thread = move_to_separate_thread(func) + self.worker.finished.connect(lambda: self.run_button.setEnabled(True)) + self.worker.finished.connect(lambda: self.root._progress_bar.hide()) + self.thread.start() + self.run_button.setEnabled(False) + self.root._progress_bar.show() From 5086d237445e281126dfff531f4b427b5d16c3c1 Mon Sep 17 00:00:00 2001 From: Jessy Lauer <30733203+jeylau@users.noreply.github.com> Date: Tue, 7 Nov 2023 15:47:22 +0100 Subject: [PATCH 10/11] Fix video adaptation (#2436) * Fix TypeError by removing trim_ends argument --- .../datasets/pose_multianimal_imgaug.py | 5 ++++- deeplabcut/pose_estimation_tensorflow/predict_supermodel.py | 5 ----- 2 files changed, 4 insertions(+), 6 deletions(-) diff --git a/deeplabcut/pose_estimation_tensorflow/datasets/pose_multianimal_imgaug.py b/deeplabcut/pose_estimation_tensorflow/datasets/pose_multianimal_imgaug.py index 5a973a39f6..edbc6894fa 100644 --- a/deeplabcut/pose_estimation_tensorflow/datasets/pose_multianimal_imgaug.py +++ b/deeplabcut/pose_estimation_tensorflow/datasets/pose_multianimal_imgaug.py @@ -383,7 +383,10 @@ def get_batch_from_video(self): for n, x, y in joints.get(j, []): kpts[j * self._n_kpts + int(n)] = x, y - joint_id = np.array(list(range(self._n_kpts)) * self._n_animals) + joint_id = [ + np.array(list(range(self._n_kpts))) + for _ in range(self._n_animals) + ] joint_ids.append(joint_id) batch_joints.append(kpts) diff --git a/deeplabcut/pose_estimation_tensorflow/predict_supermodel.py b/deeplabcut/pose_estimation_tensorflow/predict_supermodel.py index 48092ea82e..1933e013b3 100644 --- a/deeplabcut/pose_estimation_tensorflow/predict_supermodel.py +++ b/deeplabcut/pose_estimation_tensorflow/predict_supermodel.py @@ -22,7 +22,6 @@ def video_inference_superanimal( pcutoff=0.1, adapt_iterations=1000, pseudo_threshold=0.1, - trim_ends=None, ): """ Makes prediction based on a super animal model. Note right now we only support single animal video inference @@ -61,9 +60,6 @@ def video_inference_superanimal( pseudo_threshold: float, default 0.1 Video adaptation only uses predictions that are above pseudo_threshold - trim_ends: int, optional: - In cases where the beginning and ending of the videos have very messy background that impacts predictions of the model, we trim those from adaptation training - Given a list of scales for spatial pyramid, i.e. [600, 700] scale_list = range(600,800,100) @@ -107,7 +103,6 @@ def video_inference_superanimal( adapter.adaptation_training( adapt_iterations=adapt_iterations, pseudo_threshold=pseudo_threshold, - trim_ends=trim_ends, ) adapter.after_adapt_inference( pcutoff=pcutoff, From 3ba648f38abf60b440e56d8dc4bcb062bd0a04af Mon Sep 17 00:00:00 2001 From: Alexander Mathis Date: Tue, 7 Nov 2023 17:18:35 +0100 Subject: [PATCH 11/11] release updates (#2434) --- deeplabcut/version.py | 2 +- examples/test.sh | 2 +- reinstall.sh | 2 +- setup.py | 2 +- 4 files changed, 4 insertions(+), 4 deletions(-) diff --git a/deeplabcut/version.py b/deeplabcut/version.py index 969d8ec42f..bb98107cbc 100644 --- a/deeplabcut/version.py +++ b/deeplabcut/version.py @@ -9,5 +9,5 @@ # Licensed under GNU Lesser General Public License v3.0 # -__version__ = "2.3.7" +__version__ = "2.3.8" VERSION = __version__ diff --git a/examples/test.sh b/examples/test.sh index a2720ca5d8..694db69e3d 100755 --- a/examples/test.sh +++ b/examples/test.sh @@ -6,7 +6,7 @@ rm -r OUT cd .. pip uninstall deeplabcut python3 setup.py sdist bdist_wheel -pip install dist/deeplabcut-2.3.7-py3-none-any.whl +pip install dist/deeplabcut-2.3.8-py3-none-any.whl cd examples diff --git a/reinstall.sh b/reinstall.sh index de99501265..eaa0c5f733 100755 --- a/reinstall.sh +++ b/reinstall.sh @@ -1,3 +1,3 @@ pip uninstall deeplabcut python3 setup.py sdist bdist_wheel -pip install dist/deeplabcut-2.3.7-py3-none-any.whl +pip install dist/deeplabcut-2.3.8-py3-none-any.whl diff --git a/setup.py b/setup.py index 428fef611a..7f0582ba18 100644 --- a/setup.py +++ b/setup.py @@ -17,7 +17,7 @@ setuptools.setup( name="deeplabcut", - version="2.3.7", + version="2.3.8", author="A. & M. Mathis Labs", author_email="alexander@deeplabcut.org", description="Markerless pose-estimation of user-defined features with deep learning",