extending documentation #1

Author: damaggu

SwissKnife/architectures.py

@@ -33,7 +33,7 @@
     Conv2DTranspose,
     UpSampling2D,
     Reshape,
-    LeakyReLU
+    LeakyReLU,
 )
 from tensorflow.keras.models import Sequential
 
@@ -56,7 +56,10 @@ def posenet_mouse(input_shape, num_classes):
         model
     """
     recognition_model = Xception(
-        include_top=False, input_shape=input_shape, pooling="avg", weights="imagenet",
+        include_top=False,
+        input_shape=input_shape,
+        pooling="avg",
+        weights="imagenet",
     )
 
     new_input = Input(
@@ -174,7 +177,10 @@ def posenet_primate(input_shape, num_classes):  # recognition_model = DenseNet20
         num_classes:Number of classes for recognition or number of landmarks.
     """
     recognition_model = ResNet101(
-        include_top=False, input_shape=input_shape, pooling="avg", weights="imagenet",
+        include_top=False,
+        input_shape=input_shape,
+        pooling="avg",
+        weights="imagenet",
     )
 
     new_input = Input(
@@ -532,9 +538,20 @@ def classification_small(input_shape, num_classes):
 
 def dlc_model_sturman(input_shape, num_classes):
     """Model that implements behavioral classification based on Deeplabcut generated features as in Sturman et al.
-    Args:
-        input_shape:
-        num_classes:Number of behaviors to classify.
+
+    Reimplementation of the model used in the publication Sturman et al. that performs action recognition on top of pose estimation
+
+    Parameters
+    ----------
+    input_shape : keras compatible input shape (W,H,Channels)
+        keras compatible input shape (features,)
+    num_classes : int
+        Number of behaviors to classify.
+
+    Returns
+    -------
+    keras.model
+        Sturman et al. model
     """
     model = Sequential()
 
@@ -558,10 +575,21 @@ def dlc_model_sturman(input_shape, num_classes):
 
 
 def dlc_model(input_shape, num_classes):
-    """
-    Args:
-        input_shape:
-        num_classes:
+    """Model for classification on top of pose estimation.
+
+    Classification model for behavior, operating on pose estimation. This model has more free parameters than Sturman et al.
+
+    Parameters
+    ----------
+    input_shape : keras compatible input shape (W,H,Channels)
+        keras compatible input shape (features,)
+    num_classes : int
+        Number of behaviors to classify.
+
+    Returns
+    -------
+    keras.model
+        behavior (from pose estimates) model
     """
     dropout = 0.3
 
@@ -644,16 +672,19 @@ def recurrent_model_old(
 
 
 def recurrent_model_tcn(
-    recognition_model, recurrent_input_shape, classes=4,
+    recognition_model,
+    recurrent_input_shape,
+    classes=4,
 ):
-    """BehaviorNet architecture for behavioral classification based on temporal convolution architecture (TCN).
+    """Recurrent architecture for classification of temporal sequences of images based on temporal convolution architecture (TCN).
+    This architecture is used for BehaviorNet in SIPEC.
 
     Parameters
     ----------
     recognition_model : keras.model
         Pretrained recognition model that extracts features for individual frames.
-    recurrent_input_shape : np.ndarray
-        Number of classes/landmarks.
+    recurrent_input_shape : np.ndarray - (Time, Width, Height, Channels)
+        Shape of the images over time.
     classes : int
         Number of behaviors to recognise.
 
@@ -743,12 +774,24 @@ def recurrent_model_tcn(
 def recurrent_model_lstm(
     recognition_model, recurrent_input_shape, classes=4, recurrent_dropout=None
 ):
-    """
-    Args:
-        recognition_model:
-        recurrent_input_shape:
-        classes:
-        recurrent_dropout:
+    """Recurrent architecture for classification of temporal sequences of images based on LSTMs or GRUs.
+    This architecture is used for IdNet in SIPEC.
+
+    Parameters
+    ----------
+    recognition_model : keras.model
+        Pretrained recognition model that extracts features for individual frames.
+    recurrent_input_shape : np.ndarray - (Time, Width, Height, Channels)
+        Shape of the images over time.
+    classes : int
+        Number of behaviors to recognise.
+    recurrent_dropout : float
+        Recurrent dropout factor to use.
+
+    Returns
+    -------
+    keras.model
+        IdNet
     """
     input_sequences = Input(shape=recurrent_input_shape)
     sequential_model_helper = TimeDistributed(recognition_model)(input_sequences)
@@ -795,14 +838,26 @@ def recurrent_model_lstm(
     return sequential_model
 
 
+# TODO: adaptiv size
 def pretrained_recognition(model_name, input_shape, num_classes, fix_layers=True):
-    # TODO: adaptiv size
-    """
-    Args:
-        model_name:
-        input_shape:
-        num_classes:
-        fix_layers:
+    """This returns the model architecture for a model that operates on images and is pretrained with imagenet weights.
+    This architecture is used for IdNet and BehaviorNet as backbone in SIPEC and is referred to as RecognitionNet.
+
+    Parameters
+    ----------
+    model_name : keras.model
+        Name of the pretrained recognition model to use (names include: "xception, "resnet", "densenet")
+    input_shape : np.ndarray - (Time, Width, Height, Channels)
+        Shape of the images over time.
+    num_classes : int
+        Number of behaviors to recognise.
+    fix_layers : bool
+        Recurrent dropout factor to use.
+
+    Returns
+    -------
+    keras.model
+        RecognitionNet
     """
     if model_name == "xception":
         recognition_model = Xception(
@@ -912,11 +967,21 @@ def pretrained_recognition(model_name, input_shape, num_classes, fix_layers=True
 
 
 def idtracker_ai(input_shape, classes):
+    """Implementation of the idtracker.ai identification module as described in the supplementary of Romero-Ferrero et al.
+
+    Parameters
+    ----------
+    input_shape : keras compatible input shape (W,H,Channels)
+        keras compatible input shape (features,)
+    num_classes : int
+        Number of behaviors to classify..
+
+    Returns
+    -------
+    keras.model
+        idtracker.ai identification module
     """
-    Args:
-        input_shape:
-        classes:
-    """
+
     activation = "tanh"
     dropout = 0.2
     # conv model
@@ -934,7 +999,11 @@ def idtracker_ai(input_shape, classes):
     )
     model.add(Activation("relu"))
 
-    model.add(MaxPooling2D(strides=(2, 2),))
+    model.add(
+        MaxPooling2D(
+            strides=(2, 2),
+        )
+    )
 
     model.add(
         Conv2D(
@@ -947,7 +1016,11 @@ def idtracker_ai(input_shape, classes):
     )
     model.add(Activation("relu"))
 
-    model.add(MaxPooling2D(strides=(2, 2),))
+    model.add(
+        MaxPooling2D(
+            strides=(2, 2),
+        )
+    )
 
     model.add(
         Conv2D(

SwissKnife/dataloader.py

@@ -11,27 +11,21 @@
 
 
 def create_dataset(dataset, look_back=5, oneD=False):
-    # """Create a recurrent dataset from array.
-    # Args:
-    #     dataset: Numpy/List of dataset.
-    #     look_back: Number of future/past timepoints to add to current timepoint.
-    #     oneD: Boolean flag whether data is one dimensional or not.
-    # """
-    """Summary line.
-
-    Extended description of function.
+    """Create a recurrent dataset from array.
 
     Parameters
     ----------
-    arg1 : int
-        Description of arg1
-    arg2 : str
-        Description of arg2
+    dataset : np.ndarray
+        numpy array of dataset to make recurrent
+    look_back : int
+        Number of timesteps to look into the past and future.
+    oneD : bool
+        Boolean that indicates if the current dataset is one dimensional.
 
     Returns
     -------
-    bool
-        dataset
+    np.ndarray
+        recurrent dataset
     """
     dataX = []
     print("creating recurrency")
@@ -196,7 +190,6 @@ def create_recurrent_data(self, oneD=False, recurrent_labels=True):
         if recurrent_labels:
             self.create_recurrent_labels()
 
-
     def create_recurrent_data_dlc(self, recurrent_labels=True):
 
         self.dlc_train_recurrent = create_dataset(self.dlc_train, self.look_back)
@@ -209,7 +202,6 @@ def create_recurrent_data_dlc(self, recurrent_labels=True):
         if recurrent_labels:
             self.create_recurrent_labels()
 
-
     # TODO: redo all like this, i.e. gettters instead of changing data
     def expand_dims(self):
         self.x_train = np.expand_dims(self.x_train, axis=-1)
@@ -250,7 +242,7 @@ def decimate_labels(self, percentage, balanced=False):
             raise NotImplementedError
         if self.x_train_recurrent is not None:
             num_labels = int(len(self.x_train_recurrent) * percentage)
-            indices = np.arange(0, len(self.x_train_recurrent)-1)
+            indices = np.arange(0, len(self.x_train_recurrent) - 1)
             random_idxs = np.random.choice(indices, size=num_labels, replace=False)
             self.x_train = self.x_train[random_idxs]
             self.y_train = self.y_train[random_idxs]
@@ -392,7 +384,6 @@ def downscale_frames(self, factor=0.5):
             im_re.append(imresize(el, factor))
         self.x_test = np.asarray(im_re)
 
-
     def prepare_data(self, downscale=0.5, remove_behaviors=[], flatten=False):
         print("preparing data")
         self.change_dtype()

SwissKnife/full_inference.py

@@ -37,6 +37,33 @@ def full_inference(
     lookback=100,
     max_ids=4,
 ):
+    """Performs full inference on a given video using available SIPEC modules.
+
+    Parameters
+    ----------
+    videodata : np.ndarray
+        numpy array of read-in videodata.
+    results_sink : str
+        Path to where data will be saved.
+    networks : dict
+        Dictionary containing SIPEC modules to be used for full inference ("SegNet", "PoseNet", "BehaveNet", IdNet")
+    mask_matching : bool
+        Use greedy-mask-matching
+    id_matching : bool
+        Correct/smooth SIPEC:IdNet identity using identities based on temporal tracking (greedy-mask-matching)
+    mask_size : int
+        Mask size used for the cutout of animals.
+    lookback : int
+        Number of timesteps to look back into the past for id_matching.
+    max_ids : int
+        Number of maximum ids / maximum number of animals in any FOV.
+
+
+    Returns
+    -------
+    list
+        Outputs of all the provided SIPEC modules for each video frame.
+    """
     maskmatcher = MaskMatcher()
     maskmatcher.max_ids = max_ids
     classes = id_classes