Committing all app files

app.py

@@ -0,0 +1,60 @@
+import gradio as gr
+from utils.predict import predict, predict_batch
+import os
+
+inputs_list = []
+
+
+demo = gr.Blocks()
+    
+with demo:
+    
+    gr.Markdown("# **<p align='center'>ShiftViT: A Vision Transformer without Attention</p>**")
+    gr.Markdown("This space demonstrates the use of ShiftViT proposed in the paper: <a href=\"https://arxiv.org/abs/2201.10801/\">When Shift Operation Meets Vision Transformer: An Extremely Simple Alternative to Attention Mechanism</a> for image classification task.")
+    gr.Markdown("Vision Transformers (ViTs) have proven to be very useful for Computer Vision tasks. Many researchers believe that the attention layer is the main reason behind the success of ViTs.")
+    gr.Markdown("In the ShiftViT paper, the authors have tried to show that the attention mechanism may not be vital for the success of ViTs by replacing the attention operation with a shifting operation.")
+    
+    with gr.Tabs():
+        
+        with gr.TabItem("Skip Uploading!"):
+            
+            gr.Markdown("Just click *Run Model* below:")
+            with gr.Box():
+                gr.Markdown("**Prediction Probabilities** \n")
+                output_df = gr.Dataframe(headers=["image","1st_highest_probability", "2nd_highest_probability","3rd_highest_probability"],datatype=["str", "str", "str", "str"])
+                gr.Markdown("**Output Plot** \n")
+                output_plot = gr.Image(type='filepath')
+
+            gr.Markdown("**Predict**")
+            
+            with gr.Box():
+                with gr.Row():
+                    compute_button = gr.Button("Run Model")
+                
+        
+        with gr.TabItem("Upload & Predict"):
+            with gr.Box():
+                
+                with gr.Row():
+                    input_image = gr.Image(type='filepath',label="Input Image", show_label=True)
+                    output_label = gr.Label(label="Model", show_label=True)
+            
+            gr.Markdown("**Predict**")
+            
+            with gr.Box():
+                with gr.Row():
+                    submit_button = gr.Button("Submit")
+            
+            gr.Markdown("**Examples:**")
+            gr.Markdown("The model is trained to classify images belonging to the following classes:")
+
+            with gr.Column():
+                gr.Examples("examples/set2", [input_image], output_label, predict, cache_examples=True)
+        
+    
+    compute_button.click(predict_batch, inputs=input_image, outputs=[output_plot,output_df])
+    submit_button.click(predict, inputs=input_image, outputs=output_label)
+    
+    gr.Markdown('\n Author: <a href=\"https://www.linkedin.com/in/shivalika-singh/\">Shivalika Singh</a> <br> Based on this <a href=\"https://keras.io/examples/vision/shiftvit/\">Keras example</a> by <a href=\"https://twitter.com/ariG23498\">Aritra Roy Gosthipaty</a> and <a href=\"https://twitter.com/ritwik_raha\">Ritwik Raha</a> <br> Demo Powered by this <a href=\"https://huggingface.co/shivi/shiftvit/\">ShiftViT model</a>')
+    
+demo.launch()

requirements.txt

@@ -0,0 +1,3 @@
+tensorflow==2.9.1
+gradio
+tensorflow-addons

utils/constants.py

@@ -0,0 +1,32 @@
+
+class Config(object):
+    # DATA
+    batch_size = 256
+    buffer_size = batch_size * 2
+    input_shape = (32, 32, 3)
+    num_classes = 10
+
+    # AUGMENTATION
+    image_size = 48
+
+    # ARCHITECTURE
+    patch_size = 4
+    projected_dim = 96
+    num_shift_blocks_per_stages = [2, 4, 8, 2]
+    epsilon = 1e-5
+    stochastic_depth_rate = 0.2
+    mlp_dropout_rate = 0.2
+    num_div = 12
+    shift_pixel = 1
+    mlp_expand_ratio = 2
+
+    # OPTIMIZER
+    lr_start = 1e-5
+    lr_max = 1e-3
+    weight_decay = 1e-4
+
+    # TRAINING
+    epochs = 100
+
+
+class_vocab={0:"airplane",1:"automobile",2:"bird",3:"cat",4:"deer",5:"dog",6:"frog",7:"horse",8:"ship", 9:"truck"}

utils/lr_schedule.py

@@ -0,0 +1,92 @@
+import tensorflow as tf
+from tensorflow import keras
+import numpy as np
+
+"""
+Below code is taken from the [ShiftViT keras example](https://keras.io/examples/vision/shiftvit/) by Aritra Roy Gosthipaty & Ritwik Raha
+"""
+
+# Some code is taken from:
+# https://www.kaggle.com/ashusma/training-rfcx-tensorflow-tpu-effnet-b2.
+class WarmUpCosine(keras.optimizers.schedules.LearningRateSchedule):
+    """A LearningRateSchedule that uses a warmup cosine decay schedule."""
+
+    def __init__(self, lr_start, lr_max, warmup_steps, total_steps):
+        """
+        Args:
+            lr_start: The initial learning rate
+            lr_max: The maximum learning rate to which lr should increase to in
+                the warmup steps
+            warmup_steps: The number of steps for which the model warms up
+            total_steps: The total number of steps for the model training
+        """
+        super().__init__()
+        self.lr_start = lr_start
+        self.lr_max = lr_max
+        self.warmup_steps = warmup_steps
+        self.total_steps = total_steps
+        self.pi = tf.constant(np.pi)
+
+    def __call__(self, step):
+        # Check whether the total number of steps is larger than the warmup
+        # steps. If not, then throw a value error.
+        if self.total_steps < self.warmup_steps:
+            raise ValueError(
+                f"Total number of steps {self.total_steps} must be"
+                + f"larger or equal to warmup steps {self.warmup_steps}."
+            )
+
+        # `cos_annealed_lr` is a graph that increases to 1 from the initial
+        # step to the warmup step. After that this graph decays to -1 at the
+        # final step mark.
+        cos_annealed_lr = tf.cos(
+            self.pi
+            * (tf.cast(step, tf.float32) - self.warmup_steps)
+            / tf.cast(self.total_steps - self.warmup_steps, tf.float32)
+        )
+
+        # Shift the mean of the `cos_annealed_lr` graph to 1. Now the grpah goes
+        # from 0 to 2. Normalize the graph with 0.5 so that now it goes from 0
+        # to 1. With the normalized graph we scale it with `lr_max` such that
+        # it goes from 0 to `lr_max`
+        learning_rate = 0.5 * self.lr_max * (1 + cos_annealed_lr)
+
+        # Check whether warmup_steps is more than 0.
+        if self.warmup_steps > 0:
+            # Check whether lr_max is larger that lr_start. If not, throw a value
+            # error.
+            if self.lr_max < self.lr_start:
+                raise ValueError(
+                    f"lr_start {self.lr_start} must be smaller or"
+                    + f"equal to lr_max {self.lr_max}."
+                )
+
+            # Calculate the slope with which the learning rate should increase
+            # in the warumup schedule. The formula for slope is m = ((b-a)/steps)
+            slope = (self.lr_max - self.lr_start) / self.warmup_steps
+
+            # With the formula for a straight line (y = mx+c) build the warmup
+            # schedule
+            warmup_rate = slope * tf.cast(step, tf.float32) + self.lr_start
+
+            # When the current step is lesser that warmup steps, get the line
+            # graph. When the current step is greater than the warmup steps, get
+            # the scaled cos graph.
+            learning_rate = tf.where(
+                step < self.warmup_steps, warmup_rate, learning_rate
+            )
+
+        # When the current step is more that the total steps, return 0 else return
+        # the calculated graph.
+        return tf.where(
+            step > self.total_steps, 0.0, learning_rate, name="learning_rate"
+        )
+
+    def get_config(self):
+        config = {
+        "lr_start": self.lr_start,
+        "lr_max": self.lr_max, 
+        "total_steps": self.total_steps,
+        'warmup_steps': self.warmup_steps
+        }
+        return config

utils/predict.py

@@ -0,0 +1,88 @@
+import numpy as np 
+import tensorflow as tf
+from tensorflow import keras
+from huggingface_hub import from_pretrained_keras
+from .lr_schedule import WarmUpCosine
+from .constants import Config, class_vocab
+from keras.utils import load_img, img_to_array
+from tensorflow_addons.optimizers import AdamW
+import matplotlib.pyplot as plt
+import pandas as pd
+import random
+config = Config()
+
+##Load Model
+model = from_pretrained_keras("shivi/shiftvit", custom_objects={"WarmUpCosine":WarmUpCosine, "AdamW": AdamW})
+
+(x_train, y_train), (x_test, y_test) = keras.datasets.cifar10.load_data()
+
+
+AUTO = tf.data.AUTOTUNE
+
+def predict(image_path):
+    """
+    This function is used for fetching predictions corresponding to input_dataframe.
+    It outputs another dataframe containing: 
+        1. prediction probability for each class
+        2. actual expected outcome for each entry in the input dataframe
+    """
+    
+    test_image1 = load_img(image_path,target_size =(32,32))
+    test_image = img_to_array(test_image1) 
+    test_image = np.expand_dims(test_image, axis =0)
+    test_image = test_image.astype('uint8')
+    
+    
+    predict_ds = tf.data.Dataset.from_tensor_slices(test_image)
+    predict_ds = predict_ds.shuffle(config.buffer_size).batch(config.batch_size).prefetch(AUTO)
+    logits = model.predict(predict_ds) 
+    prob = tf.nn.softmax(logits)
+
+    confidences = {}
+    prob_list = prob.numpy().flatten().tolist()
+    sorted_prob = np.argsort(prob)[::-1].flatten()
+    for i in sorted_prob:
+        confidences[class_vocab[i]] = float(prob_list[i])
+    
+    return confidences
+
+
+def predict_batch(image_path):
+   
+    test_ds = tf.data.Dataset.from_tensor_slices((x_test, y_test))
+    test_ds = test_ds.batch(config.batch_size).prefetch(AUTO)
+    slice = test_ds.take(1)
+    
+    slice_pred = model.predict(slice)
+    slice_pred = tf.nn.softmax(slice_pred)
+    
+    saved_plot = "plot.jpg"
+    fig = plt.figure()
+    
+    predictions_df = pd.DataFrame()
+    num =  random.randint(0,50)
+    for images, labels in slice:
+      for i,j in zip(range(num,num+3), range(3)):
+            ax = plt.subplot(1, 3, j + 1)
+            plt.imshow(images[i].numpy().astype("uint8"))
+            output = np.argmax(slice_pred[i])
+
+            prob_list = slice_pred[i].numpy().flatten().tolist()
+            sorted_prob = np.argsort(slice_pred[i])[::-1].flatten()
+            prob_scores = {"image": "image "+ str(j), "1st_highest_probability": f"prob of {class_vocab[sorted_prob[0]]} is {round(prob_list[sorted_prob[0]] * 100,2)} %", 
+            "2nd_highest_probability": f"prob of {class_vocab[sorted_prob[1]]} is {round(prob_list[sorted_prob[1]] * 100,2)} %", 
+            "3rd_highest_probability": f"prob of {class_vocab[sorted_prob[2]]} is {round(prob_list[sorted_prob[2]] * 100,2)} %"}
+            predictions_df = predictions_df.append(prob_scores,ignore_index=True)
+            
+            plt.title(f"image {j} : {class_vocab[output]}")
+            plt.axis("off")
+            plt.savefig(saved_plot,bbox_inches='tight')
+
+    return saved_plot, predictions_df
+
+
+
+
+
+
+