Merge pull request #16 from SIPEC-Animal-Data-Analysis/docs

extending documentation #1

Author: chadhat

SwissKnife/architectures.py

@@ -54,18 +54,30 @@ def posenet(
     features=256,
     bias=False,
 ):
-    """Mouse pose estimation architecture.
-    Extended description of function.
+    """Model that implements SIPEC:PoseNet architecture.
+
+    This model uses an EfficientNet backbone and deconvolves generated features into landmarks in imagespace.
+    It operates on single images and can be used in conjuntion with SIPEC:SegNet to perform top-down pose estimation.
+
     Parameters
     ----------
-    arg1 : np.ndarray
-        Input shape for mouse pose estimation network.
-    arg2 : int
-        Number of classes/landmarks.
+    input_shape : keras compatible input shape (W,H,Channels)
+        keras compatible input shape (features,)
+    num_classes : int
+        Number of joints/landmarks to detect.
+    backbone : str
+        Backbone/feature detector to use, default is EfficientNet5. Choose smaller/bigger backbone depending on GPU memory.
+    gaussian_noise : float
+        Kernel size of gaussian noise layers to use.
+    features : int
+        Number of feature maps to generate at each level.
+    bias : bool
+        Use bias for deconvolutional layers.
+
     Returns
     -------
     keras.model
-        model
+        SIPEC:PoseNet
     """
     if backbone == "efficientnetb5":
         recognition_model = EfficientNetB5(
@@ -334,9 +346,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()
 
@@ -360,10 +383,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
 
@@ -450,14 +484,15 @@ def recurrent_model_tcn(
     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.
 
@@ -547,12 +582,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)
@@ -599,14 +646,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(
@@ -716,11 +775,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

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

@@ -40,6 +40,33 @@ def full_inference(
     mold_dimension=1024,
     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

SwissKnife/segmentation.py

@@ -9,8 +9,9 @@
 # SEGMENTATION PART
 # This code is optimized from the Mask RCNN (Waleed Abdulla, (c) 2017 Matterport, Inc.) repository
 
-#TODO: Look at the warnings and resolve them
+# TODO: Look at the warnings and resolve them
 import warnings
+
 warnings.filterwarnings("ignore")
 
 import sys
@@ -546,6 +547,7 @@ def evaluate_network(model_path, species, filter_masks=False, cv_folds=0):
 
 
 # TODO: change cv folds to None default
+# TODO: make default species
 def train_on_data_once(
     model_path,
     cv_folds=0,
@@ -557,17 +559,36 @@ def train_on_data_once(
     perform_evaluation=True,
     debug=0,
 ):
+    """Performs training for the segmentation moduel of SIPEC (SIPEC:SegNet).
 
-    """
-    Args:
-        model_path:
-        cv_folds:
-        frames_path:
-        annotations_path:
-        species:
-        fold:
-        fraction:
-        debug:
+    Parameters
+    ----------
+    model_path : str
+        Path to model, can be either where a new model should be stored or a path to an existing model to be retrained.
+    cv_folds : int
+        Number of cross_validation folds, use 0 for a normal train/test split.
+    frames_path : str
+        Path to the frames used for training.
+    annotations_path : str
+        Path to the annotations used for training.
+    species : str
+        Species to perform segmentation on (can be any species, but "mouse" or "primate" have more specialised parameters). If your species is neither "mouse" nor "primate", use "default".
+    fold : int
+        If cv_folds > 1, fold is the number of fold to be tested on.
+    fraction : float
+        Factor by which to decimate the training data points.
+    perform_evaluation : bool
+        Perform subsequent evaluation of the model
+    debug : bool
+        Debug verbosity.
+
+
+    Returns
+    -------
+    model
+        SIPEC:SegNet model
+    mean_ap
+        Mean average precision score achieved by this model
     """
     dataset_train, dataset_val = get_segmentation_data(
         frames_path=frames_path,