45 lines
1.6 KiB
Python
45 lines
1.6 KiB
Python
#图像分类
|
|
import gradio as gr
|
|
from transformers import ViTFeatureExtractor, ViTForImageClassification
|
|
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',
|
|
)
|
|
|
|
feature_extractor = ViTFeatureExtractor.from_pretrained('google/vit-base-patch16-224')
|
|
model = ViTForImageClassification.from_pretrained('google/vit-base-patch16-224')
|
|
|
|
def image_classification(image):
|
|
inputs = feature_extractor(images=image, return_tensors="pt")
|
|
outputs = model(**inputs)
|
|
logits = outputs.logits
|
|
predicted_class_idx = logits.argmax(-1).item()
|
|
return model.config.id2label[predicted_class_idx]
|
|
|
|
|
|
with gr.Blocks(theme=theme, css="footer {visibility: hidden}") as demo:
|
|
gr.Markdown("""
|
|
<div align='center' ><font size='60'>图像分类</font></div>
|
|
""")
|
|
with gr.Row():
|
|
with gr.Column():
|
|
image = gr.Image(label="图片")
|
|
with gr.Row():
|
|
button = gr.Button("提交", variant="primary")
|
|
box2 = gr.Label(num_top_classes=1, label="类别")
|
|
|
|
button.click(fn=image_classification, inputs=[image], outputs=box2)
|
|
examples = gr.Examples(examples=[['cat.jpeg'], ['dog.jpeg'], ['zebra.jpeg']], inputs=[image], label="例子")
|
|
|
|
|
|
|
|
if __name__ == "__main__":
|
|
demo.queue(concurrency_count=3).launch(server_name = "0.0.0.0")
|