MMathisLab · Aug 13, 2024 · Jul 19, 2024 · Jul 19, 2024 · Jul 25, 2024
diff --git a/deeplabcut/gui/tabs/refine_tracklets.py b/deeplabcut/gui/tabs/refine_tracklets.py
@@ -261,3 +261,5 @@ def refine_tracks(self):
             max_gap=self.max_gap_widget.value(),
         )
         self.merge_button.setEnabled(True)
+
+
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
 
@@ -295,7 +295,6 @@ def reconnect(self):
         self.lasso.connect_default_events()
         self.is_connected = True
 
-
 class TrackletVisualizer:
     def __init__(self, manager, videoname, trail_len=50):
         self.manager = manager
@@ -327,6 +326,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",
@@ -374,6 +376,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 +423,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 +449,65 @@ 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:
@@ -913,7 +977,6 @@ def filter_low_prob(cols, prob):
             df.to_hdf(output_path, key="df_with_missing", mode="w")
             df.to_csv(output_path.replace("h5", "csv"))
 
-
 def refine_tracklets(
     config,
     pickle_or_h5_file,

diff --git a/deeplabcut/utils/conversioncode.py b/deeplabcut/utils/conversioncode.py
@@ -91,6 +91,96 @@ def convertcsv2h5(config, userfeedback=True, scorer=None):
         except FileNotFoundError:
             print("Attention:", folder, "does not appear to have labeled data!")
 
+def adapt_labeled_data_to_new_project(config_path, remove_old_bodyparts = False, other_scorer = False,  userfeedback=False):
+    ''' Given the config.yaml file, this function will convert the labels of an ancient project to a new project.
+        For this, the labeled data must be in the project folder, under the labeled-data folder and with the same configuration as all deeplabcut projects.
+
+    Parameters
+    ----------
+    config_path : str
+        The path to the config.yaml file.
+    remove_old_bodyparts : bool (default = False)
+        If True, the old bodyparts that are not in the new project will be removed from the dataframe.
+    other_scorer : bool (default = False)
+        If True, the labels will be converted to the new scorer.
+    userfeedback : bool (default = True)
+        If true the user will be asked specifically for each folder in labeled-data if the containing csv shall be converted to hdf format.'''
+
+    # Load the config file
+    cfg = dlc.auxiliaryfunctions.read_config(config_path)
+
+    # Get the Project path
+    project_path = cfg['project_path']
+
+    # Get the bodyparts
+    bodyparts = cfg['multianimalbodyparts']
+    print('New Bodyparts:', bodyparts)
+
+    # Iterate over each labeled data video
+
+    # Use tqdm for a progress bar
+    for video in tqdm.tqdm(cfg['video_sets']):
+        print('Video:', video)
+
+        video_name = video.split('\\')[-1]
+        # discard the file extension
+        video_name = video_name.split('.')[0]
+        # Load the csv file
+        label_path = os.path.join(project_path, 'labeled-data', video_name)
+        csv_files = [file for file in os.listdir(label_path) if file.endswith('.csv')]
+        if not csv_files:
+            print('No csv file in the folder:', label_path)
+        else:
+            csv_path = os.path.join(label_path, csv_files[0])
+            df = pd.read_csv(csv_path, header=None)
+
+            # get the scorer
+            if other_scorer:
+                scorer = cfg['scorer']
+                # Change the scorer in the dataframe
+                df.iloc[0, 3:] = pd.Series([scorer]*len(df.columns[3:]))
+
+            else:
+                scorer = df.iloc[0, 3]
+
+            # Get the individuals
+            individuals = np.unique(df.iloc[1, 3:])
+
+            # Get the old bodyparts
+            old_bodyparts = np.unique(df.iloc[2, 3:])
+            print('Old bodyparts:', old_bodyparts)
+
+            # Get the unmber of old bodyparts
+            num_of_old_bodyparts = len(old_bodyparts)
+
+            # Bodyparts to add
+            print('Bodyparts to add:', set(bodyparts) - set(old_bodyparts))
+
+            # If a bodypart is missing, add it to the dataframe
+            for indx, bodypart in enumerate(bodyparts):
+                if bodypart not in old_bodyparts:
+                    num_of_old_bodyparts += 1
+                    for i, individual in enumerate(individuals):
+                        # create the columns for the bodypart, concatenate, the individual, the bodypart, and nan values
+                        x_column = pd.concat([pd.Series(scorer), pd.Series(individual), pd.Series(bodypart), pd.Series('x'), pd.Series(np.nan, index = df.index)], axis = 0, ignore_index = True)
+                        y_column = pd.concat([pd.Series(scorer), pd.Series(individual), pd.Series(bodypart), pd.Series('y'), pd.Series(np.nan, index = df.index)], axis = 0, ignore_index = True)
+                        # Insert the columns in the dataframe
+                        df.insert(i*2*num_of_old_bodyparts + indx*2 + 3, 'insert_' + bodypart + '_x' + individual ,x_column)
+                        df.insert(i*2*num_of_old_bodyparts + indx*2 + 4, 'insert' + bodypart + '_y' + individual, y_column)
+
+
+            # If the old bodyparts are not in the new project, remove them
+            if remove_old_bodyparts:
+                for bodypart in old_bodyparts:
+                    if bodypart not in bodyparts:
+                        df = df.drop(df.columns[df.iloc[2, :] == bodypart], axis = 1)
+
+
+            # Save the dataframe
+            df.to_csv(csv_path, index = False, header = False)
+
+    # Create/Update the h5 file
+    convertcsv2h5(config_path, userfeedback=userfeedback)
 
 def analyze_videos_converth5_to_csv(video_folder, videotype=".mp4", listofvideos=False):
     """
Original file line number	Diff line number	Diff line change
Expand Up		@@ -261,3 +261,5 @@ def refine_tracks(self):
		max_gap=self.max_gap_widget.value(),
		)
		self.merge_button.setEnabled(True)