fix #31 : generate docstrings (#36)

rtharungowda · web-flow · commit 01f24f7d85dc · 2021-02-08T18:27:57.000+05:30
diff --git a/devolearn/embryo_generator_model/embryo_generator_model.py b/devolearn/embryo_generator_model/embryo_generator_model.py
@@ -27,6 +27,13 @@
 
 class Generator(nn.Module):
     def __init__(self, ngf, nz, nc):
+        """GAN generator to generate synthetic images of embryos
+
+        Args:
+            ngf (int): size of feature maps in generator
+            nz (int): size of latent space noise (latent vector)
+            nc (int): number of color channels of the output image
+        """
         super().__init__()
         self.main = nn.Sequential(
             # input is Z, going into a convolution
@@ -65,13 +72,11 @@ def forward(self, input):
 
 class embryo_generator_model():   
     def __init__(self, device = "cpu"):
-
-        """
-        ngf = size of feature maps in generator
-        nz = size of latent space noise (latent vector)
-        nc = number of color channels of the output image
+        """Generate synthetic single or multiple images of embryos.
         Do not tweak these unless you're changing the Generator() with a new model with a different architecture. 
-    
+
+        Args:
+            device (str, optional): set to "cuda", runs operations on gpu and set to "cpu", runs operations on cpu. Defaults to "cpu".
         """
         self.device = device
         self.ngf = 128 ## generated image size 
@@ -103,19 +108,17 @@ def __init__(self, device = "cpu"):
 
 
     def generate(self, image_size = (700,500)):
-
-        """
+        """Generate one synthetic image of embryo.
+        The native size of the GAN's output is 256*256, and then it resizes the generated image to the desired size. 
         reference{
             https://github.com/DevoLearn/devolearn#generating-synthetic-images-of-embryos-with-a-pre-trained-gan
         }
-        inputs{
-            image_size <tuple> = (width,height of the generated image)
-        }
-        outputs{
-            1 channel image as an <np.array> 
-        }
-        The native size of the GAN's output is 256*256, and then it resizes the 
-        generated image to the desired size. 
+
+        Args:
+            image_size (tuple, optional): size of generated image,(width,height). Defaults to (700,500).
+
+        Returns:
+            np.array : 1 channel image 
         """
         with torch.no_grad():
             noise = torch.randn([1,128,1,1]).to(self.device)
@@ -127,21 +130,15 @@ def generate(self, image_size = (700,500)):
 
 
     def generate_n_images(self, n = 3, foldername = "generated_images", image_size = (700,500), notebook_mode = False):
-        """
+        """This is an extension of the generator.generate() function for generating multiple images at once and saving them into a folder. 
         reference{
             https://github.com/DevoLearn/devolearn#generating-synthetic-images-of-embryos-with-a-pre-trained-gan
         }
-        inputs{
-            n <int> = number of images to generate
-            foldername <str> = name of the folder where the images whould be saved. 
-            The function automatically generates a folder if it doesn't exist 
-            notebook_mode <bool> = toogle between script(False) and notebook(True), for better user interface
-        }
-        outputs{
-            None
-        }
-        
-        This is an extension of the generator.generate() function for generating multiple images at once and saving them into a folder. 
+        Args:
+            n (int, optional): number of images to generate. Defaults to 3.
+            foldername (str, optional): name of the folder where the images whould be savedThe function automatically generates a folder if it doesn't exist. Defaults to "generated_images".
+            image_size (tuple, optional): size of generated image,(width,height). Defaults to (700,500).
+            notebook_mode (bool, optional): toogle between script(False) and notebook(True), for better user interface. Defaults to False.
         """
 
         if os.path.isdir(foldername) == False:
diff --git a/devolearn/embryo_segmentor/embryo_segmentor.py b/devolearn/embryo_segmentor/embryo_segmentor.py
@@ -25,19 +25,14 @@
 """
 
 def generate_centroid_image(thresh):
-
-    """
-    used when centroid_mode == True
-
-    input{
-        thresh <np.array> = 2d numpy array that is returned from the segmentation model
-    }
-    outputs{
-        centroid_image = image containing the contours and their respective centroids 
-        centroids = list of all centroids for the given image as [(x1,y1), (x2,y2)...]
-    }
+    """Used when centroid_mode is set to True
  
+    Args:
+        thresh (np.array): 2d numpy array that is returned from the segmentation model
 
+    Returns:
+        np.array : image containing the contours and their respective centroids 
+        list : list of all centroids for the given image as [(x1,y1), (x2,y2)...]
     """
 
     thresh = cv2.blur(thresh, (5,5))
@@ -63,11 +58,11 @@ def generate_centroid_image(thresh):
 
 class embryo_segmentor():
     def __init__(self, device = "cpu"):
-        
-        """
-        Segments the c. elegans embryo from images/videos, 
+        """Segments the c. elegans embryo from images/videos, 
         depends on segmentation-models-pytorch for the model backbone
 
+        Args:
+            device (str, optional): set to "cuda", runs operations on gpu and set to "cpu", runs operations on cpu. Defaults to "cpu".
         """
         self.device = device
         self.ENCODER = 'resnet18'
@@ -110,20 +105,23 @@ def __init__(self, device = "cpu"):
 
 
     def predict(self, image_path, pred_size = (350,250), centroid_mode = False):
-
-        """
+        """Loads an image from image_path and converts it to grayscale, 
+        then passes it though the model and returns centroids of the segmented features.
         reference{
             https://github.com/DevoLearn/devolearn#segmenting-the-c-elegans-embryo
-        } 
-        inputs{
-            image_path <str> = path to image
-            pred_size <tuple> = (width,height) of the image to be returned, the default size of the model output is (256,256)
-            centroid_mode <bool> = set to true to return both the segmented image and the list of centroids 
-        }
-        outputs{
-            1 channel image as an <np.array> 
-            optional <list> containing centroids 
         }
+
+        Args:
+            image_path (str): path to image
+            pred_size (tuple, optional): size of output image,(width,height). Defaults to (350,250).
+            centroid_mode (bool, optional): set to true to return both the segmented image and the list of centroids. Defaults to False.
+
+        Returns:
+            centroid_mode set to False:
+                np.array : 1 channel image.
+            centroid_mode set to True:
+                np.array : 1 channel image,
+                list : list of centroids.
         """
 
         im = cv2.imread(image_path,0)
@@ -138,6 +136,23 @@ def predict(self, image_path, pred_size = (350,250), centroid_mode = False):
             
 
     def predict_from_video(self, video_path, pred_size = (350,250), save_folder = "preds", centroid_mode = False, notebook_mode = False):
+        """Splits a video from video_path into frames and passes the 
+        frames through the model for predictions. Saves predicted images in save_folder.
+        And optionally saves all the centroid predictions into a pandas.DataFrame. 
+
+        Args:
+            video_path (str): path to the video file.
+            pred_size (tuple, optional): size of output image,(width,height). Defaults to (350,250).
+            save_folder (str, optional): path to folder to be saved in. Defaults to "preds".
+            centroid_mode (bool, optional): set to true to return both the segmented image and the list of centroids. Defaults to False.
+            notebook_mode (bool, optional): toogle between script(False) and notebook(True), for better user interface. Defaults to False.
+
+        Returns:
+            centroid_mode set to True:
+                pd.DataFrame : containing file name and their centriods
+            centroid_mode set to False:
+                list : list containing the names of the entries in the save_folder directory 
+        """
         vidObj = cv2.VideoCapture(video_path)   
         success = 1
         images = deque()
diff --git a/devolearn/lineage_population_model/lineage_population_model.py b/devolearn/lineage_population_model/lineage_population_model.py
@@ -25,7 +25,11 @@
 
 class lineage_population_model():   
     def __init__(self, device = "cpu"):
-        
+        """Estimate lineage populations of C. elegans embroys from videos/photos and plotting predictions.
+
+        Args:
+            device (str, optional): set to "cuda", runs operations on gpu and set to "cpu", runs operations on cpu. Defaults to "cpu".
+        """
         self.device = device
         self.model = models.resnet18(pretrained = True)
         self.model.fc = nn.Linear(512, 7)  ## resize last layer
@@ -57,25 +61,20 @@ def __init__(self, device = "cpu"):
                                             ])
 
     def predict(self, image_path):
+        """Loads an image from image_path and converts it to grayscale, 
+        then passes it though the model and returns a dictionary 
+        with the scaled output (see self.scaler)
 
-        """
         reference{
             https://github.com/DevoLearn/devolearn#predicting-populations-of-cells-within-the-c-elegans-embryo
-        } 
-        input{
-            image path <str>
         }
 
-        output{
-            dictionary containing the cell population values <dict>
-        }
-
-        Loads an image from image_path and converts it to grayscale, 
-        then passes it though the model and returns a dictionary 
-        with the scaled output (see self.scaler)
+        Args:
+            image_path (str): path to image.
 
+        Returns:
+            dict: dictionary containing the cell population values
         """
-
         image = cv2.imread(image_path, 0)
         image = cv2.cvtColor(image, cv2.COLOR_GRAY2RGB)
         tensor = self.transforms(image).unsqueeze(0).to(self.device)
@@ -96,29 +95,23 @@ def predict(self, image_path):
         return pred_dict
 
     def predict_from_video(self, video_path, csv_name  = "foo.csv", save_csv = False, ignore_first_n_frames = 0, ignore_last_n_frames = 0, notebook_mode = False):
-
-        """
-        inputs{
-            video path <str> = path to video file 
-            csv_name <str> = filename to be used to save the predictions 
-            save_csv <bool> = set to True if you want to save the predictions into a CSV files
-            ignore_first_n_frames <int> = number of frames to drop in the start of the video 
-            ignore_last_n_frames <int> = number of frames to drop in the end of the video 
-            notebook_mode <bool> = toogle between script(False) and notebook(True), for better user interface
-        }
-
-
-        output{
-            DataFrame containing all the preds with the corresponding column name <pandas.DataFrame>
-        }
-        
-        Splits a video from video_path into frames and passes the 
+        """Splits a video from video_path into frames and passes the 
         frames through the model for predictions. Saves all the predictions
         into a pandas.DataFrame which can be optionally saved as a CSV file.
 
         The model was trained to make predictions upto the 
-        stage where the population of "A" lineage is 250
-
+        stage where the population of "A" lineage is 250        
+
+        Args:
+            video_path (str): path to video file
+            csv_name (str, optional): filename to be used to save the predictions. Defaults to "foo.csv".
+            save_csv (bool, optional): set to True if you want to save the predictions into a CSV files. Defaults to False.
+            ignore_first_n_frames (int, optional): number of frames to drop in the start of the video. Defaults to 0.
+            ignore_last_n_frames (int, optional): number of frames to drop in the end of the video. Defaults to 0.
+            notebook_mode (bool, optional): toogle between script(False) and notebook(True), for better user interface. Defaults to False.
+
+        Returns:
+            pandas.DataFrame : DataFrame containing all the preds with the corresponding column name
         """
         A_population_upper_limit = 250
 
@@ -180,23 +173,18 @@ def predict_from_video(self, video_path, csv_name  = "foo.csv", save_csv = False
 
         
     def create_population_plot_from_video(self, video_path, save_plot = False, plot_name = "plot.png", ignore_first_n_frames = 0, ignore_last_n_frames = 0, notebook_mode = False):
-
-        """
-        inputs{
-            video_path <str> = path to video file 
-            save_plot <bool> = set to True to save the plot as an image file 
-            plot_name <str> = filename of the plot image to be saved 
-            ignore_first_n_frames <int> = number of frames to drop in the start of the video 
-            ignore_last_n_frames <int> = number of frames to drop in the end of the video 
-            notebook_mode <bool> = toogle between script(False) and notebook(True), for better user interface
-        }
-
-        outputs{
-            plot object which can be customized further <matplotlib.pyplot>
-        }
-
-        plots all the predictions from a video into a matplotlib.pyplot 
-        
+        """Plots all the predictions from a video into a matplotlib.pyplot 
+
+        Args:
+            video_path ([type]): path to video file
+            save_plot (bool, optional): set to True to save the plot as an image file. Defaults to False.
+            plot_name (str, optional): filename of the plot image to be saved. Defaults to "plot.png".
+            ignore_first_n_frames (int, optional): number of frames to drop in the start of the video. Defaults to 0.
+            ignore_last_n_frames (int, optional): number of frames to drop in the end of the video. Defaults to 0.
+            notebook_mode (bool, optional): toogle between script(False) and notebook(True), for better user interface. Defaults to False.
+
+        Returns:
+             matplotlib.pyplot : plot object which can be customized further
         """
         df = self.predict_from_video(video_path, ignore_first_n_frames = ignore_first_n_frames, ignore_last_n_frames = ignore_last_n_frames, notebook_mode = notebook_mode)  
         
