diff --git a/app.py b/app.py index 45df5dd..b3e2e61 100644 --- a/app.py +++ b/app.py @@ -1,3 +1,34 @@ import gradio as gr +from transformers import ViTImageProcessor, ViTModel +from gradio.themes.utils import sizes +theme = gr.themes.Default(radius_size=sizes.radius_none).set( + block_label_text_color = '#4D63FF', + block_title_text_color = '#4D63FF', + button_primary_text_color = '#4D63FF', + button_primary_background_fill='#FFFFFF', + button_primary_border_color='#4D63FF', + button_primary_background_fill_hover='#EDEFFF', +) -gr.Interface.load("models/timm/vit_large_patch14_clip_224.openai_ft_in12k_in1k").launch() +processor = ViTImageProcessor.from_pretrained('google/vit-large-patch16-224-in21k') +model = ViTModel.from_pretrained('google/vit-large-patch16-224-in21k') + +def image_classification(image): + inputs = processor(images=image, return_tensors="pt") + outputs = model(**inputs) + logits = outputs.logits + predicted_label = logits.argmax(-1).item() + + return model.config.id2label[predicted_label] + +demo = gr.Interface(fn=image_classification, + inputs=gr.Image(), + outputs=gr.Label(num_top_classes=1), + title = "图像分类", + theme = theme, + examples = ['dog.jpeg']) + + +if __name__ == "__main__": + demo.queue(concurrency_count=10) + demo.launch()