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]" 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/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() 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 e8917115b1..fa458d79d4 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): @@ -395,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_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 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..79f96c539a 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( @@ -45,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() @@ -92,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) @@ -106,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: @@ -116,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. @@ -157,7 +168,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/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 d25c2e7ffb..edbc6894fa 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) @@ -162,7 +164,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 +174,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 +351,21 @@ 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) + 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 +373,32 @@ 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() + 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_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 +411,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 +501,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 +522,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 +534,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) @@ -619,7 +621,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) @@ -645,7 +647,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]) @@ -752,7 +754,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])] @@ -784,7 +786,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..1933e013b3 100644 --- a/deeplabcut/pose_estimation_tensorflow/predict_supermodel.py +++ b/deeplabcut/pose_estimation_tensorflow/predict_supermodel.py @@ -20,9 +20,8 @@ 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, ): """ 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) @@ -99,16 +95,16 @@ 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, + ) 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/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__": 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/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 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/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/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/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: 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/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/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 7cb3f20253..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", @@ -51,12 +51,14 @@ "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"], - "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 963bf2f541..a11dd148d0 100644 --- a/tests/test_dataset_augmentation.py +++ b/tests/test_dataset_augmentation.py @@ -98,12 +98,33 @@ def test_keypoint_horizontal_flip( 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) + + +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] + )[1][0] temp = keypoints_aug.reshape((3, 12, 2)) for pair in pairs: temp[:, pair] = temp[:, pair[::-1]] 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] 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)