Merge branch 'DeepLabCut:main' into main

Author: LostSunset

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

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())

deeplabcut/create_project/new.py

@@ -276,9 +276,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["corner2move2"] = (50, 50)
     cfg_file["move2corner"] = True
     cfg_file["skeleton_color"] = "black"

deeplabcut/generate_training_dataset/multiple_individuals_trainingsetmanipulation.py

@@ -272,9 +272,7 @@ def create_multianimaltraining_dataset(
     # Loading the encoder (if necessary downloading from TF)
     dlcparent_path = auxiliaryfunctions.get_deeplabcut_path()
     defaultconfigfile = os.path.join(dlcparent_path, "pose_cfg.yaml")
-    model_path = auxfun_models.check_for_weights(
-        net_type, Path(dlcparent_path)
-    )
+    model_path = auxfun_models.check_for_weights(net_type, Path(dlcparent_path))
 
     if Shuffles is None:
         Shuffles = range(1, num_shuffles + 1, 1)
@@ -425,9 +423,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,
             }

deeplabcut/generate_training_dataset/trainingsetmanipulation.py

@@ -988,9 +988,7 @@ def create_training_dataset(
             defaultconfigfile = os.path.join(dlcparent_path, "pose_cfg.yaml")
         elif posecfg_template:
             defaultconfigfile = posecfg_template
-        model_path = auxfun_models.check_for_weights(
-            net_type, Path(dlcparent_path)
-        )
+        model_path = auxfun_models.check_for_weights(net_type, Path(dlcparent_path))
 
         if Shuffles is None:
             Shuffles = range(1, num_shuffles + 1)

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

deeplabcut/gui/tabs/modelzoo.py

@@ -72,7 +72,9 @@ def _set_page(self):
 
         tooltip_label = QtWidgets.QLabel()
         tooltip_label.setPixmap(
-            QPixmap(os.path.join(BASE_DIR, "assets", "icons", "help2.png")).scaledToWidth(30)
+            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",

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):

deeplabcut/gui/window.py

@@ -59,9 +59,7 @@ def _check_for_updates(silent=True):
         _ = msg.addButton("Skip", msg.RejectRole)
         msg.exec_()
         if msg.clickedButton() is update_btn:
-            subprocess.check_call(
-                [sys.executable, "-m", *command]
-            )
+            subprocess.check_call([sys.executable, "-m", *command])
 
 
 class MainWindow(QMainWindow):

deeplabcut/modelzoo/api/superanimal_inference.py

@@ -271,7 +271,9 @@ def video_inference(
     customized_test_config="",
 ):
     if superanimal_name not in MODELOPTIONS:
-        raise ValueError(f"{superanimal_name} not available. Available ones are: {MODELOPTIONS}. If you are confident `superanimal_name` is right, try updating `dlclibrary` with `pip install -U dlclibrary`.")
+        raise ValueError(
+            f"{superanimal_name} not available. Available ones are: {MODELOPTIONS}. If you are confident `superanimal_name` is right, try updating `dlclibrary` with `pip install -U dlclibrary`."
+        )
 
     dlc_root_path = auxiliaryfunctions.get_deeplabcut_path()