talmolab · talmo · Dec 24, 2020 · Dec 16, 2020 · Dec 17, 2020 · Dec 17, 2020
diff --git a/sleap/gui/app.py b/sleap/gui/app.py
@@ -317,6 +317,12 @@ def add_submenu_choices(menu, title, options, key):
             "DeepLabCut dataset...",
             self.commands.importDLC,
         )
+        add_menu_item(
+            import_types_menu,
+            "import_dlc_folder",
+            "Multiple DeepLabCut datasets from folder...",
+            self.commands.importDLCFolder,
+        )
         add_menu_item(
             import_types_menu,
             "import_dpk",

diff --git a/sleap/gui/commands.py b/sleap/gui/commands.py
@@ -32,6 +32,7 @@ class which inherits from `AppCommand` (or a more specialized class such as
 
 from abc import ABC
 from enum import Enum
+from glob import glob
 from pathlib import PurePath
 from typing import Callable, Dict, Iterator, List, Optional, Type, Tuple
 
@@ -263,6 +264,10 @@ def importDLC(self):
         """Imports DeepLabCut datasets."""
         self.execute(ImportDeepLabCut)
 
+    def importDLCFolder(self):
+        """Imports multiple DeepLabCut datasets."""
+        self.execute(ImportDeepLabCutFolder)
+
     def importLEAP(self):
         """Imports LEAP matlab datasets."""
         self.execute(ImportLEAP)
@@ -721,6 +726,48 @@ def ask(context: "CommandContext", params: dict) -> bool:
         return True
 
 
+class ImportDeepLabCutFolder(AppCommand):
+    @staticmethod
+    def do_action(context: "CommandContext", params: dict):
+        csv_files = ImportDeepLabCutFolder.find_dlc_files_in_folder(params['folder_name'])
+        if csv_files:
+            win = MessageDialog(f"Importing {len(csv_files)} DeepLabCut datasets...", context.app)
+            merged_labels = ImportDeepLabCutFolder.import_labels_from_dlc_files(csv_files)
+            win.hide()
+
+            new_window = context.app.__class__()
+            new_window.showMaximized()
+            new_window.loadLabelsObject(labels=merged_labels)
+
+    @staticmethod
+    def ask(context: "CommandContext", params: dict) -> bool:
+        folder_name = FileDialog.openDir(
+            context.app,
+            dir=None,
+            caption="Select a folder with DeepLabCut datasets...",
+        )
+
+        if len(folder_name) == 0:
+            return False
+        params["folder_name"] = folder_name
+        return True
+
+    @staticmethod
+    def find_dlc_files_in_folder(folder_name: str) -> List[str]:
+        return glob(f"{folder_name}/*/*.csv")
+
+    @staticmethod
+    def import_labels_from_dlc_files(csv_files: List[str]) -> Labels:
+        merged_labels = None
+        for csv_file in csv_files:
+            labels = Labels.load_file(csv_file, as_format="deeplabcut")
+            if merged_labels is None:
+                merged_labels = labels
+            else:
+                merged_labels.extend_from(labels, unify=True)
+        return merged_labels
+
+
 class ImportAnalysisFile(AppCommand):
     @staticmethod
     def do_action(context: "CommandContext", params: dict):

diff --git a/tests/data/dlc_multiple_datasets/video1/dlc_dataset_1.csv b/tests/data/dlc_multiple_datasets/video1/dlc_dataset_1.csv
@@ -0,0 +1,6 @@
+scorer,Scorer,Scorer,Scorer,Scorer,Scorer,Scorer,Scorer,Scorer,Scorer,Scorer,Scorer,Scorer
+individuals,Animal1,Animal1,Animal1,Animal1,Animal1,Animal1,Animal2,Animal2,Animal2,Animal2,Animal2,Animal2
+bodyparts,A,A,B,B,C,C,A,A,B,B,C,C
+coords,x,y,x,y,x,y,x,y,x,y,x,y
+labeled-data/video/img000.jpg,0,1,2,3,4,5,6,7,8,9,10,11
+labeled-data/video/img001.jpg,12,13,,,15,16,17,18,,,20,21
diff --git a/tests/data/dlc_multiple_datasets/video1/img000.jpg b/tests/data/dlc_multiple_datasets/video1/img000.jpg
diff --git a/tests/data/dlc_multiple_datasets/video1/img001.jpg b/tests/data/dlc_multiple_datasets/video1/img001.jpg
diff --git a/tests/data/dlc_multiple_datasets/video1/img002.jpg b/tests/data/dlc_multiple_datasets/video1/img002.jpg
diff --git a/tests/data/dlc_multiple_datasets/video2/dlc_dataset_2.csv b/tests/data/dlc_multiple_datasets/video2/dlc_dataset_2.csv
@@ -0,0 +1,5 @@
+scorer,Scorer,Scorer,Scorer,Scorer,Scorer,Scorer,Scorer,Scorer,Scorer,Scorer,Scorer,Scorer
+individuals,Animal1,Animal1,Animal1,Animal1,Animal1,Animal1,Animal2,Animal2,Animal2,Animal2,Animal2,Animal2
+bodyparts,A,A,B,B,C,C,A,A,B,B,C,C
+coords,x,y,x,y,x,y,x,y,x,y,x,y
+labeled-data/video/img002.jpg,22,23,24,25,26,27,,,,,,
diff --git a/tests/data/dlc_multiple_datasets/video2/img000.jpg b/tests/data/dlc_multiple_datasets/video2/img000.jpg
diff --git a/tests/data/dlc_multiple_datasets/video2/img001.jpg b/tests/data/dlc_multiple_datasets/video2/img001.jpg
diff --git a/tests/data/dlc_multiple_datasets/video2/img002.jpg b/tests/data/dlc_multiple_datasets/video2/img002.jpg
diff --git a/tests/gui/test_commands.py b/tests/gui/test_commands.py
@@ -1,4 +1,5 @@
-from sleap.gui.commands import CommandContext
+from sleap.gui.commands import CommandContext, ImportDeepLabCutFolder
+from sleap.io.pathutils import fix_path_separator
 
 
 def test_delete_user_dialog(centered_pair_predictions):
@@ -15,3 +16,27 @@ def test_delete_user_dialog(centered_pair_predictions):
 
     # Make sure we now have user instances
     assert len(context.state["labeled_frame"].user_instances) == 2
+
+
+def test_import_labels_from_dlc_folder():
+    csv_files = ImportDeepLabCutFolder.find_dlc_files_in_folder('tests/data/dlc_multiple_datasets')
+    assert set([fix_path_separator(f) for f in csv_files]) == {
+        'tests/data/dlc_multiple_datasets/video2/dlc_dataset_2.csv',
+        'tests/data/dlc_multiple_datasets/video1/dlc_dataset_1.csv',
+    }
+
+    labels = ImportDeepLabCutFolder.import_labels_from_dlc_files(csv_files)
+
+    assert len(labels) == 3
+    assert len(labels.videos) == 2
+    assert len(labels.skeletons) == 1
+    assert len(labels.nodes) == 3
+    assert len(labels.tracks) == 0
+
+    assert set([fix_path_separator(l.video.backend.filename) for l in labels.labeled_frames]) == {
+        'tests/data/dlc_multiple_datasets/video2/img002.jpg',
+        'tests/data/dlc_multiple_datasets/video1/img000.jpg',
+        'tests/data/dlc_multiple_datasets/video1/img000.jpg',
+    }
+
+    assert set([l.frame_idx for l in labels.labeled_frames]) == {0, 0, 1}