Spaces:
Runtime error
Runtime error
| import gradio as gr | |
| import tensorflow as tf | |
| import tensorflow_hub as hub | |
| from PIL import Image | |
| import utils | |
| _RESOLUTION = 224 | |
| _MODEL_URL = "https://tfhub.dev/sayakpaul/deit_tiny_patch16_224/1" | |
| def get_model() -> tf.keras.Model: | |
| """Initiates a tf.keras.Model from TF-Hub.""" | |
| inputs = tf.keras.Input((_RESOLUTION, _RESOLUTION, 3)) | |
| hub_module = hub.KerasLayer(_MODEL_URL) | |
| logits, attention_scores_dict = hub_module( | |
| inputs | |
| ) # Second output in the tuple is a dictionary containing attention scores. | |
| return tf.keras.Model(inputs, [logits, attention_scores_dict]) | |
| _MODEL = get_model() | |
| def show_rollout(image): | |
| """Function to be called when user hits submit on the UI.""" | |
| _, preprocessed_image = utils.preprocess_image( | |
| image, "deit_tiny_patch16_224" | |
| ) | |
| _, attention_scores_dict = _MODEL.predict(preprocessed_image) | |
| result = utils.attention_rollout_map( | |
| image, attention_scores_dict, "deit_tiny_patch16_224" | |
| ) | |
| return Image.fromarray(result) | |
| title = "Generate Attention Rollout Plots" | |
| article = "Attention Rollout was proposed by [Abnar et al.](https://arxiv.org/abs/2005.00928) to quantify the information that flows through self-attention layers. In the original ViT paper ([Dosovitskiy et al.](https://arxiv.org/abs/2010.11929)), the authors use it to investigate the representations learned by ViTs. The model used in the backend is `deit_tiny_patch16_224`. For more details about it, refer [here](https://tfhub.dev/sayakpaul/collections/deit/1). DeiT was proposed by [Touvron et al.](https://arxiv.org/abs/2012.12877)" | |
| iface = gr.Interface( | |
| show_rollout, | |
| inputs=gr.inputs.Image(type="pil", label="Input Image"), | |
| outputs="image", | |
| title=title, | |
| article=article, | |
| allow_flagging="never", | |
| # examples=[["./car.jpeg", "./bulbul.jpeg"]], | |
| ) | |
| iface.launch() | |