DeepLabCut · MMathisLab · Jan 6, 2022 · Jan 5, 2022 · Jan 5, 2022 · Jan 5, 2022
diff --git a/deeplabcut/create_project/add.py b/deeplabcut/create_project/add.py
@@ -30,7 +30,7 @@ def add_new_videos(config, videos, copy_videos=False, coords=None):
     --------
     Video will be added, with cropping dimenions according to the frame dimensinos of mouse5.avi
     >>> deeplabcut.add_new_videos('/home/project/reaching-task-Tanmay-2018-08-23/config.yaml',['/data/videos/mouse5.avi'])
-    
+
     Video will be added, with cropping dimenions [0,100,0,200]
     >>> deeplabcut.add_new_videos('/home/project/reaching-task-Tanmay-2018-08-23/config.yaml',['/data/videos/mouse5.avi'],copy_videos=False,coords=[[0,100,0,200]])
 

diff --git a/deeplabcut/create_project/demo_data.py b/deeplabcut/create_project/demo_data.py
@@ -17,18 +17,22 @@
 
 def load_demo_data(config, createtrainingset=True):
     """
-  Loads the demo data. Make sure that you are in the same directory where you have downloaded or cloned the deeplabcut.
-
-  Parameter
-    ----------
-    config : string
-        Full path of the config.yaml file of the provided demo dataset as a string.
-
-    Example
-    --------
-    >>> deeplabcut.load_demo_data('config.yaml')
-    --------
-  """
+    Loads the demo data -- subset from trail-tracking data in Mathis et al. 2018. 
+    When loading, it sets paths correctly to run this project on your system
+
+    Parameter
+      ----------
+      config : string
+          Full path of the config.yaml file of the provided demo dataset as a string.
+
+      createtrainingset : bool 
+          Boolean variable indicating if a training set shall be created. 
+
+      Example
+      --------
+      >>> deeplabcut.load_demo_data('config.yaml')
+      --------
+    """
     config = Path(config).resolve()
     config = str(config)
 

diff --git a/deeplabcut/generate_training_dataset/multiple_individuals_trainingsetmanipulation.py b/deeplabcut/generate_training_dataset/multiple_individuals_trainingsetmanipulation.py
@@ -31,7 +31,10 @@
 
 
 def format_multianimal_training_data(
-    df, train_inds, project_path, n_decimals=2,
+    df,
+    train_inds,
+    project_path,
+    n_decimals=2,
 ):
     train_data = []
     nrows = df.shape[0]
@@ -283,7 +286,10 @@ def create_multianimaltraining_dataset(
 
         # Make training file!
         data = format_multianimal_training_data(
-            Data, trainIndices, cfg["project_path"], numdigits,
+            Data,
+            trainIndices,
+            cfg["project_path"],
+            numdigits,
         )
 
         if len(trainIndices) > 0:
@@ -339,7 +345,10 @@ def create_multianimaltraining_dataset(
             )
             path_test_config = str(
                 os.path.join(
-                    cfg["project_path"], Path(modelfoldername), "test", "pose_cfg.yaml",
+                    cfg["project_path"],
+                    Path(modelfoldername),
+                    "test",
+                    "pose_cfg.yaml",
                 )
             )
             path_inference_config = str(
@@ -440,7 +449,10 @@ def create_multianimaltraining_dataset(
 
 
 def convert_cropped_to_standard_dataset(
-    config_path, recreate_datasets=True, delete_crops=True, back_up=True,
+    config_path,
+    recreate_datasets=True,
+    delete_crops=True,
+    back_up=True,
 ):
     import pandas as pd
     import pickle
@@ -479,7 +491,8 @@ def convert_cropped_to_standard_dataset(
         return
 
     datasets_folder = os.path.join(
-        project_path, auxiliaryfunctions.GetTrainingSetFolder(cfg),
+        project_path,
+        auxiliaryfunctions.GetTrainingSetFolder(cfg),
     )
     df_old = pd.read_hdf(
         os.path.join(datasets_folder, "CollectedData_" + cfg["scorer"] + ".h5"),

diff --git a/deeplabcut/generate_training_dataset/trainingsetmanipulation.py b/deeplabcut/generate_training_dataset/trainingsetmanipulation.py
@@ -318,7 +318,7 @@ def check_labels(
 
 
 def boxitintoacell(joints):
-    """ Auxiliary function for creating matfile."""
+    """Auxiliary function for creating matfile."""
     outer = np.array([[None]], dtype=object)
     outer[0, 0] = np.array(joints, dtype="int64")
     return outer
@@ -460,9 +460,11 @@ def merge_annotateddatasets(cfg, trainingsetfolder_full):
 
 
 def SplitTrials(
-    trialindex, trainFraction=0.8, enforce_train_fraction=False,
+    trialindex,
+    trainFraction=0.8,
+    enforce_train_fraction=False,
 ):
-    """ Split a trial index into train and test sets. Also checks that the trainFraction is a two digit number between 0 an 1. The reason
+    """Split a trial index into train and test sets. Also checks that the trainFraction is a two digit number between 0 an 1. The reason
     is that the folders contain the trainfraction as int(100*trainFraction).
     If enforce_train_fraction is True, train and test indices are padded with -1
     such that the ratio of their lengths is exactly the desired train fraction.
@@ -487,7 +489,9 @@ def SplitTrials(
         train_indices = shuffle[: int(train_size)]
         if enforce_train_fraction and not train_size.is_integer():
             train_indices, test_indices = pad_train_test_indices(
-                train_indices, test_indices, train_fraction,
+                train_indices,
+                test_indices,
+                train_fraction,
             )
         return train_indices, test_indices
 
@@ -505,7 +509,8 @@ def pad_train_test_indices(train_inds, test_inds, train_fraction):
     min_n_train = int(round(min_length_req * train_fraction))
     min_n_test = min_length_req - min_n_train
     mult = max(
-        math.ceil(n_train_inds / min_n_train), math.ceil(n_test_inds / min_n_test),
+        math.ceil(n_train_inds / min_n_train),
+        math.ceil(n_test_inds / min_n_test),
     )
     n_train = mult * min_n_train
     n_test = mult * min_n_test
@@ -570,7 +575,8 @@ def mergeandsplit(config, trainindex=0, uniform=True):
         Data = pd.read_hdf(fn + ".h5")
     except FileNotFoundError:
         Data = merge_annotateddatasets(
-            cfg, Path(os.path.join(project_path, trainingsetfolder)),
+            cfg,
+            Path(os.path.join(project_path, trainingsetfolder)),
         )
         if Data is None:
             return [], []
@@ -582,7 +588,9 @@ def mergeandsplit(config, trainindex=0, uniform=True):
         TrainingFraction = cfg["TrainingFraction"]
         trainFraction = TrainingFraction[trainindex]
         trainIndices, testIndices = SplitTrials(
-            range(len(Data.index)), trainFraction, True,
+            range(len(Data.index)),
+            trainFraction,
+            True,
         )
     else:  # leave one folder out split
         videos = cfg["video_sets"].keys()
@@ -733,7 +741,8 @@ def create_training_dataset(
         )
 
         Data = merge_annotateddatasets(
-            cfg, Path(os.path.join(project_path, trainingsetfolder)),
+            cfg,
+            Path(os.path.join(project_path, trainingsetfolder)),
         )
         if Data is None:
             return
@@ -946,19 +955,21 @@ def create_training_dataset(
 
 
 def get_largestshuffle_index(config):
-    """ Returns the largest shuffle for all dlc-models in the current iteration."""
+    """Returns the largest shuffle for all dlc-models in the current iteration."""
     cfg = auxiliaryfunctions.read_config(config)
     project_path = cfg["project_path"]
     iterate = "iteration-" + str(cfg["iteration"])
     dlc_model_path = os.path.join(project_path, "dlc-models", iterate)
     if os.path.isdir(dlc_model_path):
         models = os.listdir(dlc_model_path)
-        # sort the models directories
+        # sort the model directories
         models.sort(key=lambda f: int("".join(filter(str.isdigit, f))))
-        # get the shuffle index
-        max_shuffle_index = int(models[-1].split("shuffle")[-1])
+
+        # get the shuffle index and offset by 1.
+        max_shuffle_index = int(models[-1].split("shuffle")[-1]) + 1
     else:
         max_shuffle_index = 0
+
     return max_shuffle_index
 
 
@@ -967,7 +978,7 @@ def create_training_model_comparison(
     trainindex=0,
     num_shuffles=1,
     net_types=["resnet_50"],
-    augmenter_types=["default"],
+    augmenter_types=["imgaug"],
     userfeedback=False,
     windows2linux=False,
 ):
@@ -1001,12 +1012,19 @@ def create_training_model_comparison(
         If this is set to false, then all requested train/test splits are created (no matter if they already exist). If you
         want to assure that previous splits etc. are not overwritten, then set this to True and you will be asked for each split.
 
+    Returns
+    ----------
+    shuffle_list: list
+        List of indices corresponding to the trainigsplits/models that were created.
+
     Example
     --------
-    >>> deeplabcut.create_training_model_comparison('/analysis/project/reaching-task/config.yaml',num_shuffles=1,net_types=['resnet_50','resnet_152'],augmenter_types=['tensorpack','deterministic'])
+    >>> shuffle_list = deeplabcut.create_training_model_comparison('/analysis/project/reaching-task/config.yaml',num_shuffles=1,net_types=['resnet_50','resnet_152'],augmenter_types=['tensorpack','deterministic'])
 
     Windows:
-    >>> deeplabcut.create_training_model_comparison('C:\\Users\\Ulf\\looming-task\\config.yaml',num_shuffles=1,net_types=['resnet_50','resnet_152'],augmenter_types=['tensorpack','deterministic'])
+    >>> shuffle_list = deeplabcut.create_training_model_comparison('C:\\Users\\Ulf\\looming-task\\config.yaml',num_shuffles=1,net_types=['resnet_50','resnet_152'],augmenter_types=['tensorpack','deterministic'])
+
+    See examples/testscript_openfielddata_augmentationcomparison.py for an example of how to use shuffle_list.
 
     --------
     """
@@ -1034,6 +1052,7 @@ def create_training_model_comparison(
 
     largestshuffleindex = get_largestshuffle_index(config)
 
+    shuffle_list = []
     for shuffle in range(num_shuffles):
         trainIndices, testIndices = mergeandsplit(
             config, trainindex=trainindex, uniform=True
@@ -1046,6 +1065,8 @@ def create_training_model_comparison(
                     + idx_net * len(augmenter_types)
                     + shuffle * len(augmenter_types) * len(net_types)
                 )
+
+                shuffle_list.append(get_max_shuffle_idx)
                 log_info = str(
                     "Shuffle index:"
                     + str(get_max_shuffle_idx)
@@ -1055,6 +1076,8 @@ def create_training_model_comparison(
                     + aug
                     + ", trainsetindex:"
                     + str(trainindex)
+                    + ", frozen shuffle ID:"
+                    + str(shuffle)
                 )
                 create_training_dataset(
                     config,
@@ -1066,3 +1089,5 @@ def create_training_model_comparison(
                     userfeedback=userfeedback,
                 )
                 logger.info(log_info)
+
+    return shuffle_list
diff --git a/deeplabcut/gui/outlier_frame_extraction_toolbox.py b/deeplabcut/gui/outlier_frame_extraction_toolbox.py
@@ -42,7 +42,6 @@ def getColorIndices(self, img, bodyparts):
         """
         Returns the colormaps ticks and . The order of ticks labels is reversed.
         """
-        #        im = io.imread(img)
         norm = mcolors.Normalize(vmin=np.min(img), vmax=np.max(img))
         ticks = np.linspace(np.min(img), np.max(img), len(bodyparts))[::-1]
         return norm, ticks

diff --git a/deeplabcut/pose_estimation_tensorflow/backbones/efficientnet_model.py b/deeplabcut/pose_estimation_tensorflow/backbones/efficientnet_model.py
@@ -386,7 +386,7 @@ def call(
 
 class Model(tf.keras.Model):
     """A class implements tf.keras.Model for MNAS-like model.
-      Reference: https://arxiv.org/abs/1807.11626
+    Reference: https://arxiv.org/abs/1807.11626
     """
 
     def __init__(self, blocks_args=None, global_params=None):