DeepLabCut · MMathisLab · Sep 2, 2024 · Jul 5, 2024 · Jul 5, 2024 · Aug 13, 2024
diff --git a/deeplabcut/benchmark/base.py b/deeplabcut/benchmark/base.py
@@ -132,10 +132,7 @@ def _validate_predictions(self, name: str, predictions: dict) -> dict:
                 "individuals were detected in those images."
             )
 
-        return {
-            img: predictions.get(img, tuple())
-            for img in test_images
-        }
+        return {img: predictions.get(img, tuple()) for img in test_images}
 
 
 @dataclasses.dataclass

diff --git a/deeplabcut/benchmark/metrics.py b/deeplabcut/benchmark/metrics.py
@@ -234,9 +234,7 @@ def calc_rmse_from_obj(
             kpts.pop(ind)
 
     test_objects = {
-        k: v
-        for k, v in eval_results_obj.items()
-        if k in gt["annotations"].keys()
+        k: v for k, v in eval_results_obj.items() if k in gt["annotations"].keys()
     }
     if len(gt["annotations"]) != len(test_objects):
         gt_images = list(gt["annotations"].keys())

diff --git a/deeplabcut/create_project/new.py b/deeplabcut/create_project/new.py
@@ -277,9 +277,9 @@ def create_new_project(
     cfg_file["x2"] = 640
     cfg_file["y1"] = 277
     cfg_file["y2"] = 624
-    cfg_file[
-        "batch_size"
-    ] = 8  # batch size during inference (video - analysis); see https://www.biorxiv.org/content/early/2018/10/30/457242
+    cfg_file["batch_size"] = (
+        8  # batch size during inference (video - analysis); see https://www.biorxiv.org/content/early/2018/10/30/457242
+    )
     cfg_file["detector_batch_size"] = 1
     cfg_file["corner2move2"] = (50, 50)
     cfg_file["move2corner"] = True

diff --git a/deeplabcut/generate_training_dataset/multiple_individuals_trainingsetmanipulation.py b/deeplabcut/generate_training_dataset/multiple_individuals_trainingsetmanipulation.py
@@ -357,9 +357,7 @@ def create_multianimaltraining_dataset(
     if engine == Engine.PYTORCH:
         model_path = dlcparent_path
     else:
-        model_path = auxfun_models.check_for_weights(
-            net_type, Path(dlcparent_path)
-        )
+        model_path = auxfun_models.check_for_weights(net_type, Path(dlcparent_path))
 
     Shuffles = validate_shuffles(cfg, Shuffles, num_shuffles, userfeedback)
 
@@ -522,9 +520,9 @@ def create_multianimaltraining_dataset(
                 "multi_step": [[1e-4, 7500], [5 * 1e-5, 12000], [1e-5, 200000]],
                 "save_iters": 10000,
                 "display_iters": 500,
-                "num_idchannel": len(cfg["individuals"])
-                if cfg.get("identity", False)
-                else 0,
+                "num_idchannel": (
+                    len(cfg["individuals"]) if cfg.get("identity", False) else 0
+                ),
                 "crop_size": list(crop_size),
                 "crop_sampling": crop_sampling,
             }

diff --git a/deeplabcut/gui/dlc_params.py b/deeplabcut/gui/dlc_params.py
@@ -13,6 +13,7 @@ class DLCParams:
         "",
         "avi",
         "mp4",
+        "mkv",
         "mov",
     ]
 

diff --git a/deeplabcut/gui/tabs/extract_frames.py b/deeplabcut/gui/tabs/extract_frames.py
@@ -95,7 +95,8 @@ def _set_page(self):
 
         self.main_layout.addWidget(
             _create_label_widget(
-                "Frame extraction from a video subset (optional for automatic extraction)", "font:bold"
+                "Frame extraction from a video subset (optional for automatic extraction)",
+                "font:bold",
             )
         )
         self.video_selection_widget = VideoSelectionWidget(self.root, self)
@@ -206,7 +207,9 @@ def extract_frames(self):
                 return
             first_video = videos[0]
             if len(videos) > 1:
-                self.root.writer.write(f"Only the first video ({first_video}) will be opened.")
+                self.root.writer.write(
+                    f"Only the first video ({first_video}) will be opened."
+                )
             video_path_in_folder = self._check_symlink(first_video)
             _ = launch_napari(str(video_path_in_folder))
             return

diff --git a/deeplabcut/gui/tracklet_toolbox.py b/deeplabcut/gui/tracklet_toolbox.py
@@ -19,7 +19,7 @@
 from deeplabcut.utils.auxfun_videos import VideoReader
 from deeplabcut.utils.auxiliaryfunctions import attempt_to_make_folder
 from matplotlib.path import Path
-from matplotlib.widgets import Slider, LassoSelector, Button, CheckButtons
+from matplotlib.widgets import Slider, LassoSelector, Button, CheckButtons, TextBox
 from PySide6.QtWidgets import QMessageBox
 from PySide6.QtCore import QMutex
 
@@ -327,6 +327,9 @@ def __init__(self, manager, videoname, trail_len=50):
 
         self.dps = []
 
+        self.swap_id1 = None
+        self.swap_id2 = None
+
     def _prepare_canvas(self, manager, fig):
         params = {
             "keymap.save": "s",
@@ -358,7 +361,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(
@@ -374,6 +377,7 @@ def _prepare_canvas(self, manager, fig):
         )
         self.vline_x = self.ax2.axvline(0, 0, 1, c="k", ls=":")
         self.vline_y = self.ax3.axvline(0, 0, 1, c="k", ls=":")
+
         custom_lines = [
             plt.Line2D([0], [0], color=self.cmap(i), lw=4)
             for i in range(len(manager.individuals))
@@ -420,10 +424,15 @@ def _prepare_canvas(self, manager, fig):
         self.ax_flag = self.fig.add_axes([0.75, 0.1, 0.05, 0.03])
         self.ax_save = self.fig.add_axes([0.80, 0.1, 0.05, 0.03])
         self.ax_help = self.fig.add_axes([0.85, 0.1, 0.05, 0.03])
+        self.ax_swap = self.fig.add_axes([0.90, 0.1, 0.05, 0.03])  # New button
+
         self.save_button = Button(self.ax_save, "Save", color="darkorange")
         self.save_button.on_clicked(self.save)
         self.help_button = Button(self.ax_help, "Help")
         self.help_button.on_clicked(self.display_help)
+        self.swap_button = Button(self.ax_swap, "Swap")  # New button
+        self.swap_button.on_clicked(self.swap_tracklets)  # Placeholder action
+
         self.drag_toggle = CheckButtons(self.ax_drag, ["Drag"])
         self.drag_toggle.on_clicked(self.toggle_draggable_points)
         self.flag_button = Button(self.ax_flag, "Flag")
@@ -441,9 +450,75 @@ def _prepare_canvas(self, manager, fig):
         self.ax1_background = self.fig.canvas.copy_from_bbox(self.ax1.bbox)
         self.fig.show()
 
+        # Create dropdowns for selecting tracklets to swap, placing them near the swap button
+        self.ax_dropdown1 = self.fig.add_axes([0.9, 0.15, 0.05, 0.03])
+        self.ax_dropdown2 = self.fig.add_axes([0.9, 0.20, 0.05, 0.03])
+        self.textbox1 = TextBox(self.ax_dropdown1, "ID 1")
+        self.textbox2 = TextBox(self.ax_dropdown2, "ID 2")
+        self.textbox1.on_submit(self.set_swap_id1)
+        self.textbox2.on_submit(self.set_swap_id2)
+
     def show(self, fig=None):
         self._prepare_canvas(self.manager, fig)
 
+    def swap_tracklets(self, event):
+        if self.swap_id1 is not None and self.swap_id2 is not None:
+
+            # Get tracklet indices for each individual
+            inds1 = [
+                k
+                for k in range(len(self.manager.tracklet2id))
+                if self.manager.tracklet2id[k] == self.swap_id1
+            ]
+            inds2 = [
+                k
+                for k in range(len(self.manager.tracklet2id))
+                if self.manager.tracklet2id[k] == self.swap_id2
+            ]
+
+            print(f"Swapping tracklets {self.swap_id1} and {self.swap_id2}")
+
+            # Frames to swap
+            frames = []
+            if len(self.cuts) == 2:
+                frames = list(range(min(self.cuts), max(self.cuts) + 1))
+            elif len(self.cuts) == 1:
+                frames = [self.cuts[0]]
+            else:
+                frames = list(range(self.curr_frame, self.manager.nframes))
+
+            # Swap the tracklets
+            for i in range(min(len(inds1), len(inds2))):
+                self.manager.swap_tracklets(inds1[i], inds2[i], frames)
+                self.display_traces()
+                self.slider.set_val(self.curr_frame)
+
+    def set_swap_id1(self, val):
+        # check that the input is a valid from the list of individuals
+        if int(val) in self.manager.tracklet2id:
+            self.swap_id1 = int(val)
+            print("ID 1 set.")
+        else:
+            print(
+                f"Invalid ID. Please select a valid ID from the list of individuals: {set(self.manager.tracklet2id)}"
+            )
+            self.swap_id1 = None
+
+    def set_swap_id2(self, val):
+        # check that the input is a valid from the list of individuals
+        if int(val) in self.manager.tracklet2id:
+            self.swap_id2 = int(val)
+            print("ID 2 set.")
+        else:
+            print(
+                f"Invalid ID. Please select a valid ID from the list of individuals: {set(self.manager.tracklet2id)}"
+            )
+            self.swap_id2 = None
+
+    def terminate(self, event):
+        plt.close(self.fig)
+        self.player.terminate()
+
     def fill_shaded_areas(self):
         self.clean_collections()
         if self.picked_pair:
@@ -587,9 +662,9 @@ def on_press(self, event):
             if len(self.cuts) > 1:
                 self.cuts.sort()
                 if self.picked_pair:
-                    self.manager.tracklet_swaps[self.picked_pair][
-                        self.cuts
-                    ] = ~self.manager.tracklet_swaps[self.picked_pair][self.cuts]
+                    self.manager.tracklet_swaps[self.picked_pair][self.cuts] = (
+                        ~self.manager.tracklet_swaps[self.picked_pair][self.cuts]
+                    )
                     self.fill_shaded_areas()
                     self.cuts = []
                     for line in self.ax_slider.lines:
@@ -807,7 +882,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/pose_estimation_3d/camera_calibration.py b/deeplabcut/pose_estimation_3d/camera_calibration.py
@@ -26,7 +26,9 @@
 matplotlib_axes_logger.setLevel("ERROR")
 
 
-def calibrate_cameras(config, cbrow=8, cbcol=6, calibrate=False, alpha=0.4, search_window_size=(11, 11)):
+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.

diff --git a/deeplabcut/pose_estimation_tensorflow/core/evaluate.py b/deeplabcut/pose_estimation_tensorflow/core/evaluate.py
@@ -868,9 +868,7 @@ def evaluate_network(
                             pose = predict.argmax_pose_predict(
                                 scmap, locref, test_pose_cfg["stride"]
                             )
-                            PredicteData[
-                                imageindex, :
-                            ] = (
+                            PredicteData[imageindex, :] = (
                                 pose.flatten()
                             )  # NOTE: thereby     cfg_test['all_joints_names'] should be same order as bodyparts!
 
@@ -972,10 +970,7 @@ def evaluate_network(
                             print("Plotting...")
                             foldername = os.path.join(
                                 str(evaluationfolder),
-                                "LabeledImages_"
-                                + DLCscorer
-                                + "_"
-                                + snapshot_name,
+                                "LabeledImages_" + DLCscorer + "_" + snapshot_name,
                             )
                             auxiliaryfunctions.attempt_to_make_folder(foldername)
                             Plotting(
@@ -998,10 +993,7 @@ def evaluate_network(
                             ).T
                             foldername = os.path.join(
                                 str(evaluationfolder),
-                                "LabeledImages_"
-                                + DLCscorer
-                                + "_"
-                                + snapshot_name,
+                                "LabeledImages_" + DLCscorer + "_" + snapshot_name,
                             )
                             if not os.path.exists(foldername):
                                 print(
@@ -1105,14 +1097,14 @@ def get_available_requested_snapshots(
 
 
 def get_snapshots_by_index(
-    idx: Union[int, str], available_snapshots: List[str],
+    idx: Union[int, str],
+    available_snapshots: List[str],
 ) -> List[str]:
     """
     Assume available_snapshots is ordered in ascending order. Returns snapshot names.
     """
-    if (
-        isinstance(idx, int)
-        and -len(available_snapshots) <= idx < len(available_snapshots)
+    if isinstance(idx, int) and -len(available_snapshots) <= idx < len(
+        available_snapshots
     ):
         return [available_snapshots[idx]]
     elif idx == "all":

diff --git a/deeplabcut/pose_estimation_tensorflow/core/evaluate_multianimal.py b/deeplabcut/pose_estimation_tensorflow/core/evaluate_multianimal.py
@@ -378,9 +378,7 @@ def evaluate_multianimal_full(
                         coords_pred = pred["coordinates"][0]
                         probs_pred = pred["confidence"]
                         for bpt, xy_gt in df.groupby(level="bodyparts"):
-                            inds_gt = np.flatnonzero(
-                                np.all(~np.isnan(xy_gt), axis=1)
-                            )
+                            inds_gt = np.flatnonzero(np.all(~np.isnan(xy_gt), axis=1))
                             n_joint = joints.index(bpt)
                             xy = coords_pred[n_joint]
                             if inds_gt.size and xy.size:
@@ -404,9 +402,9 @@ def evaluate_multianimal_full(
 
                         if plotting == "bodypart":
                             temp_xy = GT.unstack("bodyparts")[joints].values
-                            gt = temp_xy.reshape(
-                                (-1, 2, temp_xy.shape[1])
-                            ).T.swapaxes(1, 2)
+                            gt = temp_xy.reshape((-1, 2, temp_xy.shape[1])).T.swapaxes(
+                                1, 2
+                            )
                             h, w, _ = np.shape(frame)
                             fig.set_size_inches(w / 100, h / 100)
                             ax.set_xlim(0, w)
@@ -459,8 +457,7 @@ def evaluate_multianimal_full(
                     # Calculate overall prediction error
                     error = df_joint.xs("rmse", level="metrics", axis=1)
                     mask = (
-                        df_joint.xs("conf", level="metrics", axis=1)
-                        >= cfg["pcutoff"]
+                        df_joint.xs("conf", level="metrics", axis=1) >= cfg["pcutoff"]
                     )
                     error_masked = error[mask]
                     error_train = np.nanmean(error.iloc[trainIndices])
@@ -487,9 +484,7 @@ def evaluate_multianimal_full(
                             testIndices,
                         )
                         kpt_filename = DLCscorer + "-keypoint-results.csv"
-                        df_keypoint_error.to_csv(
-                            Path(evaluationfolder) / kpt_filename
-                        )
+                        df_keypoint_error.to_csv(Path(evaluationfolder) / kpt_filename)
 
                     if show_errors:
                         string = (
@@ -648,12 +643,8 @@ def evaluate_multianimal_full(
                 df.loc(axis=0)[("mAR_train", "mean")] = [
                     d[0]["mAR"] for d in results[2]
                 ]
-                df.loc(axis=0)[("mAP_test", "mean")] = [
-                    d[1]["mAP"] for d in results[2]
-                ]
-                df.loc(axis=0)[("mAR_test", "mean")] = [
-                    d[1]["mAR"] for d in results[2]
-                ]
+                df.loc(axis=0)[("mAP_test", "mean")] = [d[1]["mAP"] for d in results[2]]
+                df.loc(axis=0)[("mAR_test", "mean")] = [d[1]["mAR"] for d in results[2]]
                 with open(data_path.replace("_full.", "_map."), "wb") as file:
                     pickle.dump((df, paf_scores), file)
 

diff --git a/deeplabcut/pose_estimation_tensorflow/datasets/pose_base.py b/deeplabcut/pose_estimation_tensorflow/datasets/pose_base.py
@@ -20,12 +20,10 @@ def __init__(self, cfg):
         self.cfg = cfg
 
     @abc.abstractmethod
-    def load_dataset(self):
-        ...
+    def load_dataset(self): ...
 
     @abc.abstractmethod
-    def next_batch(self):
-        ...
+    def next_batch(self): ...
 
     def sample_scale(self):
         if self.cfg.get("deterministic", False):

diff --git 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
-Original file line number
+Diff line change
@@ Expand Up / @@ -13,6 +13,7 @@ class DLCParams: @@
             "",
             "avi",
             "mp4",
+            "mkv",
             "mov",
         ]
@@ Expand Down @@