Update README.md

README.md

@@ -7,7 +7,7 @@ tags:
 
 ## Model description
 
-ShiftViT is a variation of the Vision Transformer (ViT) where the attention operation has been swapped with a shifting operation.
+ShiftViT is a variation of the Vision Transformer (ViT) where the attention operation has been replaced with a shifting operation.
 
 ShiftViT model was proposed as part of the paper [When Shift Operation Meets Vision Transformer: An Extremely Simple Alternative to Attention Mechanism](https://arxiv.org/abs/2201.10801).
 Vision Transformers have lately become very popular for computer vision problems and a lot researchers attribute their success to the attention layers. The authors of the ShiftViT paper have tried to show via the ShiftViT model that even without the attention operation, ViTs can reach SoTA results.  
@@ -16,38 +16,44 @@ Vision Transformers have lately become very popular for computer vision problems
 
 The architecture for ShiftViT is inspired by the paper  [Swin Transformer: Hierarchical Vision Transformer using Shifted Windows](https://arxiv.org/abs/2103.14030)
 
+Here the authors propose a modular architecture with 4 stages. Each stage works on its own spatial size, creating a hierarchical architecture.
+
 | ![ShiftViT Architecture](https://i.imgur.com/CHU40HX.png) |
 | :--: |
 | Figure 1: The entire architecutre of ShiftViT.
 [Source](https://arxiv.org/abs/2201.10801) |
 
-Here the authors propose a modular architecture with 4 stages. Each stage works on its own spatial size, creating a hierarchical architecture.
-
-
 Each stage in the ShiftViT architecture comprises of a Shift Block as shown in Fig 2.
 
 | ![ShiftViT block](https://i.imgur.com/IDe35vo.gif) |
 | :--: |
 | Figure 2: From the Model to a Shift Block. |
 
-The Shift Block as shown in Fig. 3, comprises of the following:
+**The Shift Block as shown in Fig. 2, comprises of the following:**
 
-Shift Operation
-Linear Normalization
-MLP Layer: stack of densely-connected layers
+- Shift Operation
+- Linear Normalization
+- MLP Layer: stack of densely-connected layers
 
 
+**How shift operation works:**
 1. Split the channels 
 2. Select each of the first four spilts and shift and pad them in the respective directions.
 3. After shifting and padding, we concatenate the channel back.
 
  | ![Manim rendered animation for shift operation](https://i.imgur.com/PReeULP.gif) |
 | :--: |
-| Figure 4: The TensorFlow style shifting |
+| Figure 3: The TensorFlow style shifting |
 
 ## Intended uses 
 
-The ShiftViT model architecture has been designed to 
+This ShiftViT model is trained to be used for image classification task. 
+
+However, the ShiftViT architecture can be used for a variety of visual recognition tasks. 
+The authors of the [ShiftViT paper](https://arxiv.org/abs/2201.10801) tested the model on the following tasks:
+- Image Classification on ImageNet-1k dataset
+- Object Detection on COCO dataset
+- Semantic Segmentation on ADE20k dataset
 
 ## Training and evaluation data
 
@@ -76,11 +82,11 @@ No. of samples used for training and evaluation are:
 
 1. Data Preparation:
 
- - Data Augmentation: The augmentation steps include: rescaling, resizing, cropping and horizontal flipping.
+ - Data Augmentation: The augmentation steps used for data preparation include: rescaling, resizing, cropping and horizontal flipping.
  
  2. Building the ShiftViT Model:
  
- - 
+ - The steps for constructing the ShiftViT model have been covered extensively in this [Keras example](https://keras.io/examples/vision/shiftvit/)
  
  3. Model Training: