Spaces:
Runtime error
Runtime error
File size: 2,473 Bytes
310a06c 9ac731e 0f130d4 310a06c 9ac731e 310a06c 9ac731e 0f130d4 9ac731e 0f130d4 adbb278 0f130d4 9ac731e 495bdb3 9ac731e 495bdb3 310a06c 8cc2b14 310a06c 9ac731e 6f4e6e5 76d5cab 8cc2b14 9ac731e 8681243 33ca86e 8d498c1 9ac731e |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 |
# import the necessary packages
from utilities import config
from utilities import load_model
from tensorflow.keras import layers
import tensorflow as tf
import matplotlib.pyplot as plt
import math
import gradio as gr
# load the models from disk
(conv_stem, conv_trunk, conv_attn) = load_model.loader(
stem=config.IMAGENETTE_STEM_PATH,
trunk=config.IMAGENETTE_TRUNK_PATH,
attn=config.IMAGENETTE_ATTN_PATH,
)
# load labels
labels = [
'tench',
'english springer',
'cassette player',
'chain saw',
'church',
'french horn',
'garbage truck',
'gas pump',
'golf ball',
'parachute'
]
def get_results(image):
# resize the image to a 224, 224 dim
image = tf.image.convert_image_dtype(image, tf.float32)
image = tf.image.resize(image, (224, 224))
image = image[tf.newaxis, ...]
# pass through the stem
test_x = conv_stem(image)
# pass through the trunk
test_x = conv_trunk(test_x)
# pass through the attention pooling block
logits, test_viz_weights = conv_attn(test_x)
test_viz_weights = test_viz_weights[tf.newaxis, ...]
# reshape the vizualization weights
num_patches = tf.shape(test_viz_weights)[-1]
height = width = int(math.sqrt(num_patches))
test_viz_weights = layers.Reshape((height, width))(test_viz_weights)
index = 0
selected_image = image[index]
selected_weight = test_viz_weights[index]
img = plt.imshow(selected_image)
plt.imshow(
selected_weight,
cmap="inferno",
alpha=0.6,
extent=img.get_extent()
)
plt.axis("off")
prediction = tf.nn.softmax(logits, axis=-1)
prediction = prediction.numpy()[0]
return plt, {labels[i]: float(prediction[i]) for i in range(10)}
article = "<p style='text-align: center'><a href='https://arxiv.org/abs/2112.13692' target='_blank'>Augmenting Convolutional networks with attention-based aggregation</a></p> <center>Contributors: <a href='https://twitter.com/ariG23498'>Aritra Roy Gosthipaty</a>|<a href='https://twitter.com/ritwik_raha'>Ritwik Raha</a>|<a href='https://twitter.com/Cr0wley_zz'>Devjyoti Chakraborty</a></center>"
iface = gr.Interface(
fn=get_results,
title = "Patch ConvNet Demo",
description = "This space is a demo of the paper 'Augmenting Convolutional networks with attention-based aggregation' 👀",
article = article,
inputs=gr.inputs.Image(label="Input Image"),
outputs=[
gr.outputs.Image(label="Attention Map"),
gr.outputs.Label(num_top_classes=3, label="Prediction")
],
examples=[["examples/chainsaw.jpeg"], ["examples/church.jpeg"]],
).launch() |