| import tensorflow as tf | |
| import gradio as gr | |
| import gcvit | |
| from gcvit.utils import get_gradcam_model, get_gradcam_prediction | |
| def predict_fn(image, model_name): | |
| """A predict function that will be invoked by gradio.""" | |
| model = getattr(gcvit, model_name)(pretrain=True) | |
| gradcam_model = get_gradcam_model(model) | |
| preds, overlay = get_gradcam_prediction(image, gradcam_model, cmap='jet', alpha=0.4, pred_index=None) | |
| preds = {x[1]:float(x[2]) for x in preds} | |
| return [preds, overlay] | |
| demo = gr.Interface( | |
| fn=predict_fn, | |
| inputs=[ | |
| gr.inputs.Image(label="Input Image"), | |
| gr.Radio(['GCViTTiny', 'GCViTSmall', 'GCViTBase'], value='GCViTTiny', label='Model Size') | |
| ], | |
| outputs=[ | |
| gr.outputs.Label(label="Prediction"), | |
| gr.inputs.Image(label="GradCAM"), | |
| ], | |
| title="Global Context Vision Transformer (GCViT) Demo", | |
| description="Image Classification with ImageNet Pretrain Weights.", | |
| examples=[ | |
| ["example/hot_air_ballon.jpg", 'GCViTTiny'], | |
| ["example/chelsea.png", 'GCViTTiny'], | |
| ["example/german_shepherd.jpg", 'GCViTTiny'], | |
| ["example/panda.jpg", 'GCViTTiny'], | |
| ["example/jellyfish.jpg", 'GCViTTiny'], | |
| ["example/penguin.JPG", 'GCViTTiny'], | |
| ["example/bus.jpg", 'GCViTTiny'], | |
| ["example/cat_dog.JPG", 'GCViTTiny'], | |
| ], | |
| ) | |
| demo.launch() |