| | import gradio as gr |
| | from torch.nn.functional import softmax |
| | import torch |
| | from transformers import ViTFeatureExtractor |
| | from transformers import MobileViTFeatureExtractor |
| | from transformers import MobileViTForImageClassification |
| | from transformers import ViTForImageClassification |
| |
|
| |
|
| | def predict(model_type, inp): |
| |
|
| | if model_type == "ViT": |
| | model_name_or_path = './models/vit-base-garbage/' |
| | feature_extractor = ViTFeatureExtractor.from_pretrained(model_name_or_path) |
| | model = ViTForImageClassification.from_pretrained(model_name_or_path) |
| | elif model_type == "MobileViT": |
| | model_name_or_path = './models/apple/mobilevit-small-garbage/' |
| | feature_extractor = MobileViTFeatureExtractor.from_pretrained(model_name_or_path) |
| | model = MobileViTForImageClassification.from_pretrained(model_name_or_path) |
| | inputs = feature_extractor(inp, return_tensors="pt") |
| | LABELS = list(model.config.label2id.keys()) |
| | with torch.no_grad(): |
| | logits = model(**inputs) |
| | print(logits[0]) |
| | probability = torch.nn.functional.softmax(logits[0], dim=-1) |
| |
|
| | confidences = {LABELS[i]:(float(probability[0][i])) for i in range(6)} |
| | |
| | return confidences |
| |
|
| |
|
| | demo = gr.Interface(fn=predict, |
| | inputs=[gr.Dropdown(["ViT", "MobileViT"], label="Model Name", value='ViT'),gr.inputs.Image(type="pil")], |
| | outputs=gr.outputs.Label(num_top_classes=3), |
| | examples=[["ViT","paper567.jpg"],["ViT","trash105.jpg"],["ViT","plastic202.jpg"],["MobileViT","metal382.jpg"]], |
| | ) |
| | |
| | demo.launch() |
