Update MRnet.py

Author: lipi17dpatnaik

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