Spaces:
Build error
Build error
| import cv2 | |
| import json | |
| import gradio as gr | |
| import numpy as np | |
| import tensorflow as tf | |
| from backbone import create_name_vit | |
| from backbone import ClassificationModel | |
| vit_l16_384 = { | |
| "backbone_name": "vit-l/16", | |
| "backbone_params": { | |
| "image_size": 384, | |
| "representation_size": 0, | |
| "attention_dropout_rate": 0., | |
| "dropout_rate": 0., | |
| "channels": 3 | |
| }, | |
| "dropout_rate": 0., | |
| "pretrained": "./weights/vit_l16_384/model-weights" | |
| } | |
| # Init backbone | |
| backbone = create_name_vit(vit_l16_384["backbone_name"], **vit_l16_384["backbone_params"]) | |
| # Init classification model | |
| model = ClassificationModel( | |
| backbone=backbone, | |
| dropout_rate=vit_l16_384["dropout_rate"], | |
| num_classes=1000 | |
| ) | |
| # Load weights | |
| model.load_weights(vit_l16_384["pretrained"]) | |
| model.trainable = False | |
| # Load ImageNet idx to label mapping | |
| with open("assets/imagenet_1000_idx2labels.json") as f: | |
| idx_to_label = json.load(f) | |
| def resize_with_normalization(image, size=[384, 384]): | |
| image = tf.cast(image, tf.float32) | |
| image = tf.image.resize(image, size) | |
| image -= tf.constant(127.5, shape=(1, 1, 3), dtype=tf.float32) | |
| image /= tf.constant(127.5, shape=(1, 1, 3), dtype=tf.float32) | |
| image = tf.expand_dims(image, axis=0) | |
| return image | |
| def softmax_stable(x): | |
| return(np.exp(x - np.max(x)) / np.exp(x - np.max(x)).sum()) | |
| def classify_image(img, top_k): | |
| img = tf.convert_to_tensor(img) | |
| img = resize_with_normalization(img) | |
| pred_logits = model.predict(img, batch_size=1, workers=8)[0] | |
| pred_probs = softmax_stable(pred_logits) | |
| top_k_labels = pred_probs.argsort()[-top_k:][::-1] | |
| return {idx_to_label[str(idx)] : round(float(pred_probs[idx]), 4) for idx in top_k_labels} | |
| description = """ | |
| Gradio demo for <a href="https://huggingface.co/docs/transformers/main/en/model_doc/owlvit">ViT released by Kakao Lab</a>, | |
| introduced in <a href="https://arxiv.org/abs/2205.06230">Simple Open-Vocabulary Object Detection | |
| with Vision Transformers</a>. | |
| \n\nYou can use OWL-ViT to query images with text descriptions of any object. | |
| To use it, simply upload an image and enter comma separated text descriptions of objects you want to query the image for. You | |
| can also use the score threshold slider to set a threshold to filter out low probability predictions. | |
| \n\nOWL-ViT is trained on text templates, | |
| hence you can get better predictions by querying the image with text templates used in training the original model: *"photo of a star-spangled banner"*, | |
| *"image of a shoe"*. Refer to the <a href="https://arxiv.org/abs/2103.00020">CLIP</a> paper to see the full list of text templates used to augment the training data. | |
| """ | |
| demo = gr.Interface( | |
| classify_image, | |
| inputs=[gr.Image(), gr.Slider(0, 1000, value=5)], | |
| outputs=gr.outputs.Label(), | |
| title="Image Classification with Kakao Brain ViT", | |
| #description=description, | |
| ) | |
| demo.launch() |