Merge pull request #501 from lipi17dpatnaik/patch-1

Add additional comments

Author: vikasgupta-github

MRNet-Single-Model/dataset/dataset.py

@@ -19,52 +19,58 @@ class MRData():
 
     def __init__(self,task = 'acl', train = True, transform = None, weights = None):
         """Initialize the dataset
-
         Args:
             plane : along which plane to load the data
             task : for which task to load the labels
             train : whether to load the train or val data
             transform : which transforms to apply
             weights (Tensor) : Give wieghted loss to postive class eg. `weights=torch.tensor([2.223])`
         """
+        # Define the three planes to use
         self.planes=['axial', 'coronal', 'sagittal']
+        # Initialize the records as None
         self.records = None
         # an empty dictionary
         self.image_path={}
 
+        # If we are in training loop
         if train:
+            # Read data about patient records
             self.records = pd.read_csv('./images/train-{}.csv'.format(task),header=None, names=['id', 'label'])
 
-            '''
-            self.image_path[<plane>]= dictionary {<plane>: path to folder containing
-                                                                image for that plane}
-            '''
             for plane in self.planes:
+                # For each plane, specify the image path
                 self.image_path[plane] = './images/train/{}/'.format(plane)
         else:
+            # If we are in testing loop
+            # don't use any transformation
             transform = None
+            # Read testing/validation data (patients records)
             self.records = pd.read_csv('./images/valid-{}.csv'.format(task),header=None, names=['id', 'label'])
-            '''
-            self.image_path[<plane>]= dictionary {<plane>: path to folder containing
-                                                                image for that plane}
-            '''
+            
             for plane in self.planes:
+                # Read path of images for each plane
                 self.image_path[plane] = './images/valid/{}/'.format(plane)
 
-        
+        # Initialize the transformation to apply on images
         self.transform = transform 
 
+        # Append 0s to the patient record id
         self.records['id'] = self.records['id'].map(
             lambda i: '0' * (4 - len(str(i))) + str(i))
         # empty dictionary
-        self.paths={}    
+        self.paths={}
         for plane in self.planes:
+            # Get paths of numpy data files for each plane
             self.paths[plane] = [self.image_path[plane] + filename +
                           '.npy' for filename in self.records['id'].tolist()]
 
+        # Convert labels from Pandas Series to a list
         self.labels = self.records['label'].tolist()
 
+        # Total positive cases
         pos = sum(self.labels)
+        # Total negative cases
         neg = len(self.labels) - pos
 
         # Find the wieghts of pos and neg classes
@@ -90,53 +96,75 @@ def __getitem__(self, index):
         img_raw = {}
 
         for plane in self.planes:
+            # Load raw image data for each plane
             img_raw[plane] = np.load(self.paths[plane][index])
+            # Resize the image loaded in the previous step
             img_raw[plane] = self._resize_image(img_raw[plane])
 
         label = self.labels[index]
+        # Convert label to 0 and 1
         if label == 1:
             label = torch.FloatTensor([1])
         elif label == 0:
             label = torch.FloatTensor([0])
 
+        # Return a list of three images for three planes and the label of the record
         return [img_raw[plane] for plane in self.planes], label
 
     def _resize_image(self, image):
         """Resize the image to `(3,224,224)` and apply 
         transforms if possible.
         """
         # Resize the image
+        # Calculate extra padding present in the image
+        # which needs to be removed
         pad = int((image.shape[2] - INPUT_DIM)/2)
+        # This is equivalent to center cropping the image
         image = image[:,pad:-pad,pad:-pad]
+        # Normalize the image by subtracting it by mean and dividing by standard
+        # deviation
         image = (image-np.min(image))/(np.max(image)-np.min(image))*MAX_PIXEL_VAL
         image = (image - MEAN) / STDDEV
-
+        
+        # If the transformation is not None
         if self.transform:
+            # Transform the image based on the specified transformation
             image = self.transform(image)
         else:
+            # Else, just stack the image with itself in order to match the required
+            # dimensions
             image = np.stack((image,)*3, axis=1)
-        
+        # Convert the image to a FloatTensor and return it
         image = torch.FloatTensor(image)
         return image
 
 def load_data(task : str):
 
     # Define the Augmentation here only
     augments = Compose([
+        # Convert the image to Tensor
         transforms.Lambda(lambda x: torch.Tensor(x)),
+        # Randomly rotate the image with an angle
+        # between -25 degrees to 25 degrees
         RandomRotate(25),
+        # Randomly translate the image by 11% of 
+        # image height and width
         RandomTranslate([0.11, 0.11]),
+        # Randomly flip the image
         RandomFlip(),
+        # Change the order of image channels
         transforms.Lambda(lambda x: x.repeat(3, 1, 1, 1).permute(1, 0, 2, 3)),
     ])
 
     print('Loading Train Dataset of {} task...'.format(task))
+    # Load training dataset
     train_data = MRData(task, train=True, transform=augments)
     train_loader = data.DataLoader(
         train_data, batch_size=1, num_workers=11, shuffle=True
     )
 
     print('Loading Validation Dataset of {} task...'.format(task))
+    # Load validation dataset
     val_data = MRData(task, train=False)
     val_loader = data.DataLoader(
         val_data, batch_size=1, num_workers=11, shuffle=False

MRNet-Single-Model/models/MRnet.py

@@ -7,19 +7,31 @@ class MRnet(nn.Module):
     """MRnet uses pretrained resnet50 as a backbone to extract features
     """
 
-    def __init__(self): # add conf file
-
+    def __init__(self):
+        """This function will be used to initialize the 
+        MRnet instance."""
+        # Initialize nn.Module instance
         super(MRnet,self).__init__()
 
-        # init three backbones for three axis
+        # Initialize three backbones for three axis
+        # All the three axes will use pretrained AlexNet model
+        # The models will be used for extracting features from
+        # the input images
         self.axial = models.alexnet(pretrained=True).features
         self.coronal = models.alexnet(pretrained=True).features
         self.saggital = models.alexnet(pretrained=True).features
-
+        
+        # Initialize 2D Adaptive Average Pooling layers
+        # The pooling layers will reduce the size of
+        # feature maps extracted from the previous axes
         self.pool_axial = nn.AdaptiveAvgPool2d(1)
         self.pool_coronal = nn.AdaptiveAvgPool2d(1)
         self.pool_saggital = nn.AdaptiveAvgPool2d(1)
-
+        
+        # Initialize a sequential neural network with
+        # a single fully connected linear layer
+        # The network will output the probability of
+        # having a particular disease
         self.fc = nn.Sequential(
             nn.Linear(in_features=3*256,out_features=1)
         )
@@ -33,24 +45,37 @@ def forward(self,x):
         # squeeze the first dimension as there
         # is only one patient in each batch
         images = [torch.squeeze(img, dim=0) for img in x]
-
+        
+        # Extract features across each of the three plane
+        # using the three pre-trained AlexNet models defined earlier
         image1 = self.axial(images[0])
         image2 = self.coronal(images[1])
         image3 = self.saggital(images[2])
 
+        # Convert the image dimesnsions from [slices, 256, 1, 1] to
+        # [slices,256]
         image1 = self.pool_axial(image1).view(image1.size(0), -1)
         image2 = self.pool_coronal(image2).view(image2.size(0), -1)
         image3 = self.pool_saggital(image3).view(image3.size(0), -1)
 
+        # Find maximum value in each slice
+        # This will reduce the dimensions of image to [1,256]
+        # This is done in order to keep only the most prevalent
+        # features for each slice
         image1 = torch.max(image1,dim=0,keepdim=True)[0]
         image2 = torch.max(image2,dim=0,keepdim=True)[0]
         image3 = torch.max(image3,dim=0,keepdim=True)[0]
 
+        # Stack the 3 images together to create the output
+        # of size [1, 256*3]
         output = torch.cat([image1,image2,image3], dim=1)
 
+        # Feed the output to the sequential network created earlier
+        # The network will return a probability of having a specific
+        # disease
         output = self.fc(output)
         return output
 
     def _load_wieghts(self):
         """load pretrained weights"""
-        pass
+        pass

MRNet-Single-Model/utils/utils.py

@@ -15,38 +15,49 @@ def _evaluate_model(model, val_loader, criterion, epoch, num_epochs, writer, cur
 
     # Set to eval mode
     model.eval()
-
+    # List of probabilities obtained from the model
     y_probs = []
+    # List of groundtruth labels
     y_gt = []
+    # List of losses obtained
     losses = []
 
+    # Iterate over the validation dataset
     for i, (images, label) in enumerate(val_loader):
-
+        # If GPU is available, load the images and label
+        # on GPU
         if torch.cuda.is_available():
             images = [image.cuda() for image in images]
             label = label.cuda()
 
+        # Obtain the model output by passing the images as input
         output = model(images)
-
+        # Evaluate the loss by comparing the output and groundtruth label
         loss = criterion(output, label)
-
+        # Add loss to the list of losses
         loss_value = loss.item()
         losses.append(loss_value)
-
+        # Find probability for each class by applying
+        # sigmoid function on model output
         probas = torch.sigmoid(output)
-
+        # Add the groundtruth to the list of groundtruths
         y_gt.append(int(label.item()))
+        # Add predicted probability to the list
         y_probs.append(probas.item())
 
         try:
+            # Evaluate area under ROC curve based on the groundtruth label
+            # and predicted probability
             auc = metrics.roc_auc_score(y_gt, y_probs)
         except:
+            # Default area under ROC curve
             auc = 0.5
-
+        # Add information to the writer about validation loss and Area under ROC curve
         writer.add_scalar('Val/Loss', loss_value, epoch * len(val_loader) + i)
         writer.add_scalar('Val/AUC', auc, epoch * len(val_loader) + i)
 
         if (i % log_every == 0) & (i > 0):
+            # Display the information about average validation loss and area under ROC curve
             print('''[Epoch: {0} / {1} | Batch : {2} / {3} ]| Avg Val Loss {4} | Val AUC : {5} | lr : {6}'''.
                   format(
                       epoch + 1,
@@ -58,9 +69,9 @@ def _evaluate_model(model, val_loader, criterion, epoch, num_epochs, writer, cur
                       current_lr
                   )
                   )
-
+    # Add information to the writer about total epochs and Area under ROC curve
     writer.add_scalar('Val/AUC_epoch', auc, epoch + i)
-
+    # Find mean area under ROC curve and validation loss
     val_loss_epoch = np.round(np.mean(losses), 4)
     val_auc_epoch = np.round(auc, 4)
 
@@ -71,41 +82,62 @@ def _train_model(model, train_loader, epoch, num_epochs, optimizer, criterion, w
     # Set to train mode
     model.train()
 
+    # Initialize the predicted probabilities
     y_probs = []
+    # Initialize the groundtruth labels
     y_gt = []
+    # Initialize the loss between the groundtruth label
+    # and the predicted probability
     losses = []
 
+    # Iterate over the training dataset
     for i, (images, label) in enumerate(train_loader):
+        # Reset the gradient by zeroing it
         optimizer.zero_grad()
-
+        
+        # If GPU is available, transfer the images and label
+        # to the GPU
         if torch.cuda.is_available():
             images = [image.cuda() for image in images]
             label = label.cuda()
 
+        # Obtain the prediction using the model
         output = model(images)
 
+        # Evaluate the loss by comparing the prediction
+        # and groundtruth label
         loss = criterion(output, label)
+        # Perform a backward propagation
         loss.backward()
+        # Modify the weights based on the error gradient
         optimizer.step()
 
+        # Add current loss to the list of losses
         loss_value = loss.item()
         losses.append(loss_value)
 
+        # Find probabilities from output using sigmoid function
         probas = torch.sigmoid(output)
 
+        # Add current groundtruth label to the list of groundtruths
         y_gt.append(int(label.item()))
+        # Add current probabilities to the list of probabilities
         y_probs.append(probas.item())
 
         try:
+            # Try finding the area under ROC curve
             auc = metrics.roc_auc_score(y_gt, y_probs)
         except:
+            # Use default value of area under ROC curve as 0.5
             auc = 0.5
-
+        
+        # Add information to the writer about training loss and Area under ROC curve
         writer.add_scalar('Train/Loss', loss_value,
                           epoch * len(train_loader) + i)
         writer.add_scalar('Train/AUC', auc, epoch * len(train_loader) + i)
 
         if (i % log_every == 0) & (i > 0):
+            # Display the information about average training loss and area under ROC curve
             print('''[Epoch: {0} / {1} | Batch : {2} / {3} ]| Avg Train Loss {4} | Train AUC : {5} | lr : {6}'''.
                   format(
                       epoch + 1,
@@ -117,9 +149,10 @@ def _train_model(model, train_loader, epoch, num_epochs, optimizer, criterion, w
                       current_lr
                   )
                   )
-
+    # Add information to the writer about total epochs and Area under ROC curve
     writer.add_scalar('Train/AUC_epoch', auc, epoch + i)
 
+    # Find mean area under ROC curve and training loss
     train_loss_epoch = np.round(np.mean(losses), 4)
     train_auc_epoch = np.round(auc, 4)